Fix some whitespace issues in discrete.py

Author: aseyboldt

pymc3/distributions/discrete.py

@@ -632,11 +632,13 @@ def logp(self, value):
         psi = self.psi
         theta = self.theta
 
-        logp_val = tt.switch(tt.gt(value, 0),
-                     tt.log(psi) + self.pois.logp(value),
-                     logaddexp(tt.log1p(-psi), tt.log(psi) - theta))
+        logp_val = tt.switch(
+            tt.gt(value, 0),
+            tt.log(psi) + self.pois.logp(value),
+            logaddexp(tt.log1p(-psi), tt.log(psi) - theta))
 
-        return bound(logp_val,
+        return bound(
+            logp_val,
             0 <= value,
             0 <= psi, psi <= 1,
             0 <= theta)
@@ -700,11 +702,13 @@ def logp(self, value):
         p = self.p
         n = self.n
 
-        logp_val = tt.switch(tt.gt(value, 0),
-                 tt.log(psi) + self.bin.logp(value),
-                 logaddexp(tt.log1p(-psi), tt.log(psi) + n * tt.log1p(-p)))
+        logp_val = tt.switch(
+            tt.gt(value, 0),
+            tt.log(psi) + self.bin.logp(value),
+            logaddexp(tt.log1p(-psi), tt.log(psi) + n * tt.log1p(-p)))
 
-        return bound(logp_val,
+        return bound(
+            logp_val,
             0 <= value, value <= n,
             0 <= psi, psi <= 1,
             0 <= p, p <= 1)
@@ -715,10 +719,14 @@ def _repr_latex_(self, name=None, dist=None):
         n = dist.n
         p = dist.p
         psi = dist.psi
-        return r'${} \sim \text{{ZeroInflatedBinomial}}(\mathit{{n}}={}, \mathit{{p}}={}, \mathit{{psi}}={})$'.format(name,
-                                                get_variable_name(n),
-                                                get_variable_name(p),
-                                                get_variable_name(psi))
+
+        name_n = get_variable_name(n)
+        name_p = get_variable_name(p)
+        name_psi = get_variable_name(psi)
+        return (r'${} \sim \text{{ZeroInflatedBinomial}}'
+                r'(\mathit{{n}}={}, \mathit{{p}}={}, '
+                r'\mathit{{psi}}={})$'
+                .format(name, name_n, name_p, name_psi))
 
 
 class ZeroInflatedNegativeBinomial(Discrete):
@@ -731,10 +739,18 @@ class ZeroInflatedNegativeBinomial(Discrete):
 
     .. math::
 
-       f(x \mid \psi, \mu, \alpha) = \left\{ \begin{array}{l}
-            (1-\psi) + \psi \left (\frac{\alpha}{\alpha+\mu} \right) ^\alpha, \text{if } x = 0 \\
-            \psi \frac{\Gamma(x+\alpha)}{x! \Gamma(\alpha)} \left (\frac{\alpha}{\mu+\alpha} \right)^\alpha \left( \frac{\mu}{\mu+\alpha} \right)^x, \text{if } x=1,2,3,\ldots
-            \end{array} \right.
+       f(x \mid \psi, \mu, \alpha) = \left\{
+         \begin{array}{l}
+           (1-\psi) + \psi \left (
+             \frac{\alpha}{\alpha+\mu}
+           \right) ^\alpha, \text{if } x = 0 \\
+           \psi \frac{\Gamma(x+\alpha)}{x! \Gamma(\alpha)} \left (
+             \frac{\alpha}{\mu+\alpha}
+           \right)^\alpha \left(
+             \frac{\mu}{\mu+\alpha}
+           \right)^x, \text{if } x=1,2,3,\ldots
+         \end{array}
+       \right.
 
     ========  ==========================
     Support   :math:`x \in \mathbb{N}_0`
@@ -776,22 +792,33 @@ def logp(self, value):
         mu = self.mu
         psi = self.psi
 
-        logp_val = tt.switch(tt.gt(value, 0),
-                     tt.log(psi) + self.nb.logp(value),
-                     logaddexp(tt.log1p(-psi), tt.log(psi) + alpha * (tt.log(alpha) - tt.log(alpha + mu))))
+        logp_other = tt.log(psi) + self.nb.logp(value)
+        logp_0 = logaddexp(
+            tt.log1p(-psi),
+            tt.log(psi) + alpha * (tt.log(alpha) - tt.log(alpha + mu)))
+
+        logp_val = tt.switch(
+            tt.gt(value, 0),
+            logp_other,
+            logp_0)
 
-        return bound(logp_val,
-                    0 <= value,
-                    0 <= psi, psi <= 1,
-                    mu > 0, alpha > 0)
+        return bound(
+            logp_val,
+            0 <= value,
+            0 <= psi, psi <= 1,
+            mu > 0, alpha > 0)
 
     def _repr_latex_(self, name=None, dist=None):
         if dist is None:
             dist = self
         mu = dist.mu
         alpha = dist.alpha
         psi = dist.psi
-        return r'${} \sim \text{{ZeroInflatedNegativeBinomial}}(\mathit{{mu}}={}, \mathit{{alpha}}={}, \mathit{{psi}}={})$'.format(name,
-                                                get_variable_name(mu),
-                                                get_variable_name(alpha),
-                                                get_variable_name(psi))
+
+        name_mu = get_variable_name(mu)
+        name_alpha = get_variable_name(alpha)
+        name_psi = get_variable_name(psi)
+        return (r'${} \sim \text{{ZeroInflatedNegativeBinomial}}'
+                r'(\mathit{{mu}}={}, \mathit{{alpha}}={}, '
+                r'\mathit{{psi}}={})$'
+                .format(name, name_mu, name_alpha, name_psi))