add underscore escape fix to distribution repr as well as model repr,…

… fixed testing to expect underscores in LaTeX representation to be escaped
pymc-devs · Sep 14, 2024 · 67383c5 · 67383c5
1 parent 753082f
commit 67383c5
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 10 deletions.
diff --git a/pymc/printing.py b/pymc/printing.py
@@ -58,6 +58,7 @@ def str_for_dist(
     if "latex" in formatting:
         if print_name is not None:
             print_name = r"\text{" + _latex_escape(print_name.strip("$")) + "}"
+            print_name = _format_underscore(print_name)
 
         op_name = (
             dist.owner.op._print_name[1]
@@ -307,8 +308,15 @@ def _format_underscore(variable: str) -> str:
     """
     if "_" not in variable:
         return variable
+
     inds = [i for i, ltr in enumerate(variable) if ltr == "_"]
-    for i, ind in enumerate(inds):
-        ind = ind + i
-        variable = variable[:ind] + "\\" + variable[ind:]
+    var_len_original = len(variable)
+    var_len = None
+    for ind in inds:
+        if var_len:
+            if var_len != var_len_original:
+                ind = ind + (var_len - var_len_original)
+        if variable[ind - 1 : ind] != "\\":
+            variable = variable[:ind] + "\\" + variable[ind:]
+            var_len = len(variable)
     return variable
diff --git a/tests/test_printing.py b/tests/test_printing.py
@@ -173,15 +173,15 @@ def setup_class(self):
                 r"$\text{mu} \sim \operatorname{Deterministic}(f(\text{beta},~\text{alpha}))$",
                 r"$\text{beta} \sim \operatorname{Normal}(0,~10)$",
                 r"$\text{Z} \sim \operatorname{MultivariateNormal}(f(),~f())$",
-                r"$\text{nb_with_p_n} \sim \operatorname{NegativeBinomial}(10,~\text{nbp})$",
+                r"$\text{nb\_with\_p\_n} \sim \operatorname{NegativeBinomial}(10,~\text{nbp})$",
                 r"$\text{zip} \sim \operatorname{MarginalMixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5))$",
                 r"$\text{w} \sim \operatorname{Dirichlet}(\text{<constant>})$",
                 (
-                    r"$\text{nested_mix} \sim \operatorname{MarginalMixture}(\text{w},"
+                    r"$\text{nested\_mix} \sim \operatorname{MarginalMixture}(\text{w},"
                     r"~\operatorname{MarginalMixture}(f(),~\operatorname{DiracDelta}(0),~\operatorname{Poisson}(5)),"
                     r"~\operatorname{Censored}(\operatorname{Bernoulli}(0.5),~-1,~1))$"
                 ),
-                r"$\text{Y_obs} \sim \operatorname{Normal}(\text{mu},~\text{sigma})$",
+                r"$\text{Y\_obs} \sim \operatorname{Normal}(\text{mu},~\text{sigma})$",
                 r"$\text{pot} \sim \operatorname{Potential}(f(\text{beta},~\text{alpha}))$",
                 r"$\text{pred} \sim \operatorname{Deterministic}(f(\text{<normal>}))",
             ],
@@ -191,11 +191,11 @@ def setup_class(self):
                 r"$\text{mu} \sim \operatorname{Deterministic}$",
                 r"$\text{beta} \sim \operatorname{Normal}$",
                 r"$\text{Z} \sim \operatorname{MultivariateNormal}$",
-                r"$\text{nb_with_p_n} \sim \operatorname{NegativeBinomial}$",
+                r"$\text{nb\_with\_p\_n} \sim \operatorname{NegativeBinomial}$",
                 r"$\text{zip} \sim \operatorname{MarginalMixture}$",
                 r"$\text{w} \sim \operatorname{Dirichlet}$",
-                r"$\text{nested_mix} \sim \operatorname{MarginalMixture}$",
-                r"$\text{Y_obs} \sim \operatorname{Normal}$",
+                r"$\text{nested\_mix} \sim \operatorname{MarginalMixture}$",
+                r"$\text{Y\_obs} \sim \operatorname{Normal}$",
                 r"$\text{pot} \sim \operatorname{Potential}$",
                 r"$\text{pred} \sim \operatorname{Deterministic}",
             ],
@@ -258,7 +258,7 @@ def test_model_latex_repr_three_levels_model():
         "$$",
         "\\begin{array}{rcl}",
         "\\text{mu} &\\sim & \\operatorname{Normal}(0,~5)\\\\\\text{sigma} &\\sim & "
-        "\\operatorname{HalfCauchy}(0,~2.5)\\\\\\text{censored_normal} &\\sim & "
+        "\\operatorname{HalfCauchy}(0,~2.5)\\\\\\text{censored\\_normal} &\\sim & "
         "\\operatorname{Censored}(\\operatorname{Normal}(\\text{mu},~\\text{sigma}),~-2,~2)",
         "\\end{array}",
         "$$",