update if conditions, formatting defaults and tests

+ latex formatting should be detected by `if "latex" in formatting` to catch both format options
+ all latex reprs except for an entire model default to "latex_with_params"
+ tests now cover cases with and without params

Author: michaelosthege

pymc3/distributions/bart.py

@@ -246,7 +246,7 @@ def _str_repr(self, name=None, dist=None, formatting="plain"):
         alpha = self.alpha
         m = self.m
 
-        if formatting == "latex":
+        if "latex" in formatting:
             return f"$\\text{{{name}}} \\sim  \\text{{BART}}(\\text{{alpha = }}\\text{{{alpha}}}, \\text{{m = }}\\text{{{m}}})$"
         else:
             return f"{name} ~ BART(alpha = {alpha}, m = {m})"

pymc3/distributions/bound.py

@@ -157,13 +157,13 @@ def _distr_name_for_repr(self):
 
     def _str_repr(self, **kwargs):
         distr_repr = self._wrapped._str_repr(**{**kwargs, "dist": self._wrapped})
-        if "formatting" in kwargs and kwargs["formatting"] == "latex":
+        if "formatting" in kwargs and "latex" in kwargs["formatting"]:
             distr_repr = distr_repr[distr_repr.index(r" \sim") + 6 :]
         else:
             distr_repr = distr_repr[distr_repr.index(" ~") + 3 :]
         self_repr = super()._str_repr(**kwargs)
 
-        if "formatting" in kwargs and kwargs["formatting"] == "latex":
+        if "formatting" in kwargs and "latex" in kwargs["formatting"]:
             return self_repr + " -- " + distr_repr
         else:
             return self_repr + "-" + distr_repr

pymc3/distributions/distribution.py

@@ -216,9 +216,9 @@ def __str__(self, **kwargs):
         except:
             return super().__str__()
 
-    def _repr_latex_(self, **kwargs):
+    def _repr_latex_(self, *, formatting="latex_with_params", **kwargs):
         """Magic method name for IPython to use for LaTeX formatting."""
-        return self._str_repr(formatting="latex_with_params", **kwargs)
+        return self._str_repr(formatting=formatting, **kwargs)
 
     def logp_nojac(self, *args, **kwargs):
         """Return the logp, but do not include a jacobian term for transforms.

pymc3/distributions/simulator.py

@@ -126,7 +126,7 @@ def _str_repr(self, name=None, dist=None, formatting="plain"):
         sum_stat = self.sum_stat.__name__ if hasattr(self.sum_stat, "__call__") else self.sum_stat
         distance = getattr(self.distance, "__name__", self.distance.__class__.__name__)
 
-        if formatting == "latex":
+        if "latex" in formatting:
             return f"$\\text{{{name}}} \\sim  \\text{{Simulator}}(\\text{{{function}}}({params}), \\text{{{distance}}}, \\text{{{sum_stat}}})$"
         else:
             return f"{name} ~ Simulator({function}({params}), {distance}, {sum_stat})"

pymc3/model.py

@@ -65,7 +65,7 @@ def __rmatmul__(self, other):
 
     def _str_repr(self, name=None, dist=None, formatting="plain"):
         if getattr(self, "distribution", None) is None:
-            if formatting == "latex":
+            if "latex" in formatting:
                 return None
             else:
                 return super().__str__()
@@ -76,8 +76,8 @@ def _str_repr(self, name=None, dist=None, formatting="plain"):
             dist = self.distribution
         return self.distribution._str_repr(name=name, dist=dist, formatting=formatting)
 
-    def _repr_latex_(self, **kwargs):
-        return self._str_repr(formatting="latex", **kwargs)
+    def _repr_latex_(self, *, formatting="latex_with_params", **kwargs):
+        return self._str_repr(formatting=formatting, **kwargs)
 
     def __str__(self, **kwargs):
         try:
@@ -1375,8 +1375,8 @@ def check_test_point(self, test_point=None, round_vals=2):
     def _str_repr(self, formatting="plain", **kwargs):
         all_rv = itertools.chain(self.unobserved_RVs, self.observed_RVs)
 
-        if formatting == "latex":
-            rv_reprs = [rv.__latex__() for rv in all_rv]
+        if "latex" in formatting:
+            rv_reprs = [rv.__latex__(formatting=formatting) for rv in all_rv]
             rv_reprs = [
                 rv_repr.replace(r"\sim", r"&\sim &").strip("$")
                 for rv_repr in rv_reprs
@@ -1407,8 +1407,8 @@ def _str_repr(self, formatting="plain", **kwargs):
     def __str__(self, **kwargs):
         return self._str_repr(formatting="plain", **kwargs)
 
-    def _repr_latex_(self, **kwargs):
-        return self._str_repr(formatting="latex", **kwargs)
+    def _repr_latex_(self, *, formatting="latex", **kwargs):
+        return self._str_repr(formatting=formatting, **kwargs)
 
     __latex__ = _repr_latex_
 
@@ -1893,8 +1893,8 @@ def _str_repr(self, formatting="plain"):
                 return f"{self.name} ~ Deterministic({args})"
             return f"{self.name} ~ Deterministic"
 
-    def _repr_latex_(self):
-        return self._str_repr(formatting="latex_with_params")
+    def _repr_latex_(self, *, formatting="latex_with_params", **kwargs):
+        return self._str_repr(formatting=formatting)
 
     __latex__ = _repr_latex_
 

pymc3/tests/test_distributions.py

@@ -1800,58 +1800,83 @@ def setup_class(self):
             Y_obs = Normal("Y_obs", mu=mu, sigma=sigma, observed=Y)
 
         self.distributions = [alpha, sigma, mu, b, Z, Y_obs, bound_var]
-        self.expected_latex = (
-            r"$\text{alpha} \sim \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
-            r"$\text{sigma} \sim \text{HalfNormal}(\mathit{sigma}=1.0)$",
-            r"$\text{mu} \sim \text{Deterministic}(\text{alpha},~\text{Constant},~\text{beta})$",
-            r"$\text{beta} \sim \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
-            r"$\text{Z} \sim \text{MvNormal}(\mathit{mu}=array,~\mathit{chol_cov}=array)$",
-            r"$\text{Y_obs} \sim \text{Normal}(\mathit{mu}=\text{mu},~\mathit{sigma}=f(\text{sigma}))$",
-            r"$\text{bound_var} \sim \text{Bound}(\mathit{lower}=1.0,~\mathit{upper}=\text{None})$ -- \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
-            r"$\text{kron_normal} \sim \text{KroneckerNormal}(\mathit{mu}=array)$",
-            r"$\text{mat_normal} \sim \text{MatrixNormal}(\mathit{mu}=array,~\mathit{rowcov}=array,~\mathit{colchol_cov}=array)$",
-        )
-        self.expected_str = (
-            r"alpha ~ Normal(mu=0.0, sigma=10.0)",
-            r"sigma ~ HalfNormal(sigma=1.0)",
-            r"mu ~ Deterministic(alpha, Constant, beta)",
-            r"beta ~ Normal(mu=0.0, sigma=10.0)",
-            r"Z ~ MvNormal(mu=array, chol_cov=array)",
-            r"Y_obs ~ Normal(mu=mu, sigma=f(sigma))",
-            r"bound_var ~ Bound(lower=1.0, upper=None)-Normal(mu=0.0, sigma=10.0)",
-            r"kron_normal ~ KroneckerNormal(mu=array)",
-            r"mat_normal ~ MatrixNormal(mu=array, rowcov=array, colchol_cov=array)",
-        )
+        self.expected = {
+            "latex": (
+                r"$\text{alpha} \sim \text{Normal}$",
+                r"$\text{sigma} \sim \text{HalfNormal}$",
+                r"$\text{mu} \sim \text{Deterministic}$",
+                r"$\text{beta} \sim \text{Normal}$",
+                r"$\text{Z} \sim \text{MvNormal}$",
+                r"$\text{Y_obs} \sim \text{Normal}$",
+                r"$\text{bound_var} \sim \text{Bound}$ -- \text{Normal}$",
+                r"$\text{kron_normal} \sim \text{KroneckerNormal}$",
+                r"$\text{mat_normal} \sim \text{MatrixNormal}$",
+            ),
+            "plain": (
+                r"alpha ~ Normal",
+                r"sigma ~ HalfNormal",
+                r"mu ~ Deterministic",
+                r"beta ~ Normal",
+                r"Z ~ MvNormal",
+                r"Y_obs ~ Normal",
+                r"bound_var ~ Bound-Normal",
+                r"kron_normal ~ KroneckerNormal",
+                r"mat_normal ~ MatrixNormal",
+            ),
+            "latex_with_params": (
+                r"$\text{alpha} \sim \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
+                r"$\text{sigma} \sim \text{HalfNormal}(\mathit{sigma}=1.0)$",
+                r"$\text{mu} \sim \text{Deterministic}(\text{alpha},~\text{Constant},~\text{beta})$",
+                r"$\text{beta} \sim \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
+                r"$\text{Z} \sim \text{MvNormal}(\mathit{mu}=array,~\mathit{chol_cov}=array)$",
+                r"$\text{Y_obs} \sim \text{Normal}(\mathit{mu}=\text{mu},~\mathit{sigma}=f(\text{sigma}))$",
+                r"$\text{bound_var} \sim \text{Bound}(\mathit{lower}=1.0,~\mathit{upper}=\text{None})$ -- \text{Normal}(\mathit{mu}=0.0,~\mathit{sigma}=10.0)$",
+                r"$\text{kron_normal} \sim \text{KroneckerNormal}(\mathit{mu}=array)$",
+                r"$\text{mat_normal} \sim \text{MatrixNormal}(\mathit{mu}=array,~\mathit{rowcov}=array,~\mathit{colchol_cov}=array)$",
+            ),
+            "plain_with_params": (
+                r"alpha ~ Normal(mu=0.0, sigma=10.0)",
+                r"sigma ~ HalfNormal(sigma=1.0)",
+                r"mu ~ Deterministic(alpha, Constant, beta)",
+                r"beta ~ Normal(mu=0.0, sigma=10.0)",
+                r"Z ~ MvNormal(mu=array, chol_cov=array)",
+                r"Y_obs ~ Normal(mu=mu, sigma=f(sigma))",
+                r"bound_var ~ Bound(lower=1.0, upper=None)-Normal(mu=0.0, sigma=10.0)",
+                r"kron_normal ~ KroneckerNormal(mu=array)",
+                r"mat_normal ~ MatrixNormal(mu=array, rowcov=array, colchol_cov=array)",
+            ),
+        }
 
     def test__repr_latex_(self):
-        for distribution, tex in zip(self.distributions, self.expected_latex):
+        for distribution, tex in zip(self.distributions, self.expected["latex_with_params"]):
             assert distribution._repr_latex_() == tex
 
         model_tex = self.model._repr_latex_()
 
-        for tex in self.expected_latex:  # make sure each variable is in the model
+        # make sure each variable is in the model
+        for tex in self.expected["latex"]:
             for segment in tex.strip("$").split(r"\sim"):
                 assert segment in model_tex
 
     def test___latex__(self):
-        for distribution, tex in zip(self.distributions, self.expected_latex):
+        for distribution, tex in zip(self.distributions, self.expected["latex_with_params"]):
             assert distribution._repr_latex_() == distribution.__latex__()
         assert self.model._repr_latex_() == self.model.__latex__()
 
     def test___str__(self):
-        for distribution, str_repr in zip(self.distributions, self.expected_str):
+        for distribution, str_repr in zip(self.distributions, self.expected["plain"]):
             assert distribution.__str__() == str_repr
 
         model_str = self.model.__str__()
-        for str_repr in self.expected_str:
+        for str_repr in self.expected["plain"]:
             assert str_repr in model_str
 
     def test_str(self):
-        for distribution, str_repr in zip(self.distributions, self.expected_str):
+        for distribution, str_repr in zip(self.distributions, self.expected["plain"]):
             assert str(distribution) == str_repr
 
         model_str = str(self.model)
-        for str_repr in self.expected_str:
+        for str_repr in self.expected["plain"]:
             assert str_repr in model_str
 
 

pymc3/util.py

@@ -137,7 +137,7 @@ def get_repr_for_variable(variable, formatting="plain"):
                     for item in variable.get_parents()[0].inputs
                 ]
                 # do not escape_latex these, since it is not idempotent
-                if formatting == "latex":
+                if "latex" in formatting:
                     return "f({args})".format(
                         args=",~".join([n for n in names if isinstance(n, str)])
                     )
@@ -152,7 +152,7 @@ def get_repr_for_variable(variable, formatting="plain"):
             return value.item()
         return "array"
 
-    if formatting == "latex":
+    if "latex" in formatting:
         return fr"\text{{{name}}}"
     else:
         return name