Enable ruff to format code in docstrings

Author: Armavica

pymc/backends/__init__.py

@@ -23,7 +23,7 @@
 Values can be accessed in a few ways. The easiest way is to index the
 backend object with a variable or variable name.
 
-    >>> trace['x']  # or trace.x or trace[x]
+    >>> trace["x"]  # or trace.x or trace[x]
 
 The call will return the sampling values of `x`, with the values for
 all chains concatenated. (For a single call to `sample`, the number of
@@ -32,18 +32,18 @@
 To discard the first N values of each chain, slicing syntax can be
 used.
 
-    >>> trace['x', 1000:]
+    >>> trace["x", 1000:]
 
 The `get_values` method offers more control over which values are
 returned. The call below will discard the first 1000 iterations
 from each chain and keep the values for each chain as separate arrays.
 
-    >>> trace.get_values('x', burn=1000, combine=False)
+    >>> trace.get_values("x", burn=1000, combine=False)
 
 The `chains` parameter of `get_values` can be used to limit the chains
 that are retrieved.
 
-    >>> trace.get_values('x', burn=1000, chains=[0, 2])
+    >>> trace.get_values("x", burn=1000, chains=[0, 2])
 
 MultiTrace objects also support slicing. For example, the following
 call would return a new trace object without the first 1000 sampling

pymc/data.py

@@ -390,16 +390,16 @@ def Data(
     >>> observed_data = [mu + np.random.randn(20) for mu in true_mu]
 
     >>> with pm.Model() as model:
-    ...     data = pm.Data('data', observed_data[0])
-    ...     mu = pm.Normal('mu', 0, 10)
-    ...     pm.Normal('y', mu=mu, sigma=1, observed=data)
+    ...     data = pm.Data("data", observed_data[0])
+    ...     mu = pm.Normal("mu", 0, 10)
+    ...     pm.Normal("y", mu=mu, sigma=1, observed=data)
 
     >>> # Generate one trace for each dataset
     >>> idatas = []
     >>> for data_vals in observed_data:
     ...     with model:
     ...         # Switch out the observed dataset
-    ...         model.set_data('data', data_vals)
+    ...         model.set_data("data", data_vals)
     ...         idatas.append(pm.sample())
     """
     if coords is None:

pymc/distributions/continuous.py

@@ -488,10 +488,10 @@ class Normal(Continuous):
     .. code-block:: python
 
         with pm.Model():
-            x = pm.Normal('x', mu=0, sigma=10)
+            x = pm.Normal("x", mu=0, sigma=10)
 
         with pm.Model():
-            x = pm.Normal('x', mu=0, tau=1/23)
+            x = pm.Normal("x", mu=0, tau=1 / 23)
     """
 
     rv_op = normal
@@ -636,13 +636,13 @@ class TruncatedNormal(BoundedContinuous):
     .. code-block:: python
 
         with pm.Model():
-            x = pm.TruncatedNormal('x', mu=0, sigma=10, lower=0)
+            x = pm.TruncatedNormal("x", mu=0, sigma=10, lower=0)
 
         with pm.Model():
-            x = pm.TruncatedNormal('x', mu=0, sigma=10, upper=1)
+            x = pm.TruncatedNormal("x", mu=0, sigma=10, upper=1)
 
         with pm.Model():
-            x = pm.TruncatedNormal('x', mu=0, sigma=10, lower=0, upper=1)
+            x = pm.TruncatedNormal("x", mu=0, sigma=10, lower=0, upper=1)
 
     """
 
@@ -817,10 +817,10 @@ class HalfNormal(PositiveContinuous):
     .. code-block:: python
 
         with pm.Model():
-            x = pm.HalfNormal('x', sigma=10)
+            x = pm.HalfNormal("x", sigma=10)
 
         with pm.Model():
-            x = pm.HalfNormal('x', tau=1/15)
+            x = pm.HalfNormal("x", tau=1 / 15)
     """
 
     rv_op = halfnormal
@@ -1711,10 +1711,10 @@ class LogNormal(PositiveContinuous):
 
         # Example to show that we pass in only ``sigma`` or ``tau`` but not both.
         with pm.Model():
-            x = pm.LogNormal('x', mu=2, sigma=30)
+            x = pm.LogNormal("x", mu=2, sigma=30)
 
         with pm.Model():
-            x = pm.LogNormal('x', mu=2, tau=1/100)
+            x = pm.LogNormal("x", mu=2, tau=1 / 100)
     """
 
     rv_op = lognormal
@@ -1828,10 +1828,10 @@ class StudentT(Continuous):
     .. code-block:: python
 
         with pm.Model():
-            x = pm.StudentT('x', nu=15, mu=0, sigma=10)
+            x = pm.StudentT("x", nu=15, mu=0, sigma=10)
 
         with pm.Model():
-            x = pm.StudentT('x', nu=15, mu=0, lam=1/23)
+            x = pm.StudentT("x", nu=15, mu=0, lam=1 / 23)
     """
 
     rv_op = t
@@ -2802,10 +2802,10 @@ class HalfStudentT(PositiveContinuous):
 
         # Only pass in one of lam or sigma, but not both.
         with pm.Model():
-            x = pm.HalfStudentT('x', sigma=10, nu=10)
+            x = pm.HalfStudentT("x", sigma=10, nu=10)
 
         with pm.Model():
-            x = pm.HalfStudentT('x', lam=4, nu=10)
+            x = pm.HalfStudentT("x", lam=4, nu=10)
     """
 
     rv_type = HalfStudentTRV
@@ -4104,9 +4104,9 @@ class PolyaGamma(PositiveContinuous):
 
         rng = np.random.default_rng()
         with pm.Model():
-            x = pm.PolyaGamma('x', h=1, z=5.5)
+            x = pm.PolyaGamma("x", h=1, z=5.5)
         with pm.Model():
-            x = pm.PolyaGamma('x', h=25, z=-2.3, rng=rng, size=(100, 5))
+            x = pm.PolyaGamma("x", h=25, z=-2.3, rng=rng, size=(100, 5))
 
     References
     ----------

pymc/distributions/custom.py

@@ -571,13 +571,15 @@ class CustomDist:
         import pymc as pm
         from pytensor.tensor import TensorVariable
 
+
         def logp(value: TensorVariable, mu: TensorVariable) -> TensorVariable:
-            return -(value - mu)**2
+            return -((value - mu) ** 2)
+
 
         with pm.Model():
-            mu = pm.Normal('mu',0,1)
+            mu = pm.Normal("mu", 0, 1)
             pm.CustomDist(
-                'custom_dist',
+                "custom_dist",
                 mu,
                 logp=logp,
                 observed=np.random.randn(100),
@@ -596,20 +598,23 @@ def logp(value: TensorVariable, mu: TensorVariable) -> TensorVariable:
         import pymc as pm
         from pytensor.tensor import TensorVariable
 
+
         def logp(value: TensorVariable, mu: TensorVariable) -> TensorVariable:
-            return -(value - mu)**2
+            return -((value - mu) ** 2)
+
 
         def random(
             mu: np.ndarray | float,
             rng: Optional[np.random.Generator] = None,
-            size : Optional[Tuple[int]]=None,
-        ) -> np.ndarray | float :
+            size: Optional[Tuple[int]] = None,
+        ) -> np.ndarray | float:
             return rng.normal(loc=mu, scale=1, size=size)
 
+
         with pm.Model():
-            mu = pm.Normal('mu', 0 , 1)
+            mu = pm.Normal("mu", 0, 1)
             pm.CustomDist(
-                'custom_dist',
+                "custom_dist",
                 mu,
                 logp=logp,
                 random=random,
@@ -629,13 +634,15 @@ def random(
         import pymc as pm
         from pytensor.tensor import TensorVariable
 
+
         def dist(
             lam: TensorVariable,
             shift: TensorVariable,
             size: TensorVariable,
         ) -> TensorVariable:
             return pm.Exponential.dist(lam, size=size) + shift
 
+
         with pm.Model() as m:
             lam = pm.HalfNormal("lam")
             shift = -1

pymc/distributions/mixture.py

@@ -194,10 +194,10 @@ class Mixture(Distribution):
 
         # Mixture of 2 Poisson variables
         with pm.Model() as model:
-            w = pm.Dirichlet('w', a=np.array([1, 1]))  # 2 mixture weights
+            w = pm.Dirichlet("w", a=np.array([1, 1]))  # 2 mixture weights
 
-            lam1 = pm.Exponential('lam1', lam=1)
-            lam2 = pm.Exponential('lam2', lam=1)
+            lam1 = pm.Exponential("lam1", lam=1)
+            lam2 = pm.Exponential("lam2", lam=1)
 
             # As we just need the logp, rather than add a RV to the model, we need to call `.dist()`
             # These two forms are equivalent, but the second benefits from vectorization
@@ -208,14 +208,14 @@ class Mixture(Distribution):
             # `shape=(2,)` indicates 2 mixture components
             components = pm.Poisson.dist(mu=pm.math.stack([lam1, lam2]), shape=(2,))
 
-            like = pm.Mixture('like', w=w, comp_dists=components, observed=data)
+            like = pm.Mixture("like", w=w, comp_dists=components, observed=data)
 
 
     .. code-block:: python
 
         # Mixture of Normal and StudentT variables
         with pm.Model() as model:
-            w = pm.Dirichlet('w', a=np.array([1, 1]))  # 2 mixture weights
+            w = pm.Dirichlet("w", a=np.array([1, 1]))  # 2 mixture weights
 
             mu = pm.Normal("mu", 0, 1)
 
@@ -224,7 +224,7 @@ class Mixture(Distribution):
                 pm.StudentT.dist(nu=4, mu=mu, sigma=1),
             ]
 
-            like = pm.Mixture('like', w=w, comp_dists=components, observed=data)
+            like = pm.Mixture("like", w=w, comp_dists=components, observed=data)
 
 
     .. code-block:: python
@@ -233,10 +233,10 @@ class Mixture(Distribution):
         with pm.Model() as model:
             # w is a stack of 5 independent size 3 weight vectors
             # If shape was `(3,)`, the weights would be shared across the 5 replication dimensions
-            w = pm.Dirichlet('w', a=np.ones(3), shape=(5, 3))
+            w = pm.Dirichlet("w", a=np.ones(3), shape=(5, 3))
 
             # Each of the 3 mixture components has an independent mean
-            mu = pm.Normal('mu', mu=np.arange(3), sigma=1, shape=3)
+            mu = pm.Normal("mu", mu=np.arange(3), sigma=1, shape=3)
 
             # These two forms are equivalent, but the second benefits from vectorization
             components = [
@@ -249,14 +249,14 @@ class Mixture(Distribution):
             # The mixture is an array of 5 elements
             # Each element can be thought of as an independent scalar mixture of 3
             # components with different means
-            like = pm.Mixture('like', w=w, comp_dists=components, observed=data)
+            like = pm.Mixture("like", w=w, comp_dists=components, observed=data)
 
 
     .. code-block:: python
 
         # Mixture of 2 Dirichlet variables
         with pm.Model() as model:
-            w = pm.Dirichlet('w', a=np.ones(2))  # 2 mixture weights
+            w = pm.Dirichlet("w", a=np.ones(2))  # 2 mixture weights
 
             # These two forms are equivalent, but the second benefits from vectorization
             components = [
@@ -267,7 +267,7 @@ class Mixture(Distribution):
 
             # The mixture is an array of 3 elements
             # Each element comes from only one of the two core Dirichlet components
-            like = pm.Mixture('like', w=w, comp_dists=components, observed=data)
+            like = pm.Mixture("like", w=w, comp_dists=components, observed=data)
     """
 
     rv_type = MarginalMixtureRV

pymc/distributions/multivariate.py

@@ -230,34 +230,36 @@ class MvNormal(Continuous):
     Define a multivariate normal variable for a given covariance
     matrix::
 
-        cov = np.array([[1., 0.5], [0.5, 2]])
+        cov = np.array([[1.0, 0.5], [0.5, 2]])
         mu = np.zeros(2)
-        vals = pm.MvNormal('vals', mu=mu, cov=cov, shape=(5, 2))
+        vals = pm.MvNormal("vals", mu=mu, cov=cov, shape=(5, 2))
 
     Most of the time it is preferable to specify the cholesky
     factor of the covariance instead. For example, we could
     fit a multivariate outcome like this (see the docstring
     of `LKJCholeskyCov` for more information about this)::
 
         mu = np.zeros(3)
-        true_cov = np.array([[1.0, 0.5, 0.1],
-                             [0.5, 2.0, 0.2],
-                             [0.1, 0.2, 1.0]])
+        true_cov = np.array(
+            [
+                [1.0, 0.5, 0.1],
+                [0.5, 2.0, 0.2],
+                [0.1, 0.2, 1.0],
+            ],
+        )
         data = np.random.multivariate_normal(mu, true_cov, 10)
 
         sd_dist = pm.Exponential.dist(1.0, shape=3)
-        chol, corr, stds = pm.LKJCholeskyCov('chol_cov', n=3, eta=2,
-            sd_dist=sd_dist, compute_corr=True)
-        vals = pm.MvNormal('vals', mu=mu, chol=chol, observed=data)
+        chol, corr, stds = pm.LKJCholeskyCov("chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True)
+        vals = pm.MvNormal("vals", mu=mu, chol=chol, observed=data)
 
     For unobserved values it can be better to use a non-centered
     parametrization::
 
         sd_dist = pm.Exponential.dist(1.0, shape=3)
-        chol, _, _ = pm.LKJCholeskyCov('chol_cov', n=3, eta=2,
-            sd_dist=sd_dist, compute_corr=True)
-        vals_raw = pm.Normal('vals_raw', mu=0, sigma=1, shape=(5, 3))
-        vals = pm.Deterministic('vals', pt.dot(chol, vals_raw.T).T)
+        chol, _, _ = pm.LKJCholeskyCov("chol_cov", n=3, eta=2, sd_dist=sd_dist, compute_corr=True)
+        vals_raw = pm.Normal("vals_raw", mu=0, sigma=1, shape=(5, 3))
+        vals = pm.Deterministic("vals", pt.dot(chol, vals_raw.T).T)
     """
 
     rv_op = multivariate_normal
@@ -1806,13 +1808,12 @@ class MatrixNormal(Continuous):
     Define a matrixvariate normal variable for given row and column covariance
     matrices::
 
-        colcov = np.array([[1., 0.5], [0.5, 2]])
+        colcov = np.array([[1.0, 0.5], [0.5, 2]])
         rowcov = np.array([[1, 0, 0], [0, 4, 0], [0, 0, 16]])
         m = rowcov.shape[0]
         n = colcov.shape[0]
         mu = np.zeros((m, n))
-        vals = pm.MatrixNormal('vals', mu=mu, colcov=colcov,
-                               rowcov=rowcov)
+        vals = pm.MatrixNormal("vals", mu=mu, colcov=colcov, rowcov=rowcov)
 
     Above, the ith row in vals has a variance that is scaled by 4^i.
     Alternatively, row or column cholesky matrices could be substituted for
@@ -2418,23 +2419,25 @@ class ICAR(Continuous):
         # 4x4 adjacency matrix
         # arranged in a square lattice
 
-        W = np.array([
-            [0,1,0,1],
-            [1,0,1,0],
-            [0,1,0,1],
-            [1,0,1,0]
-        ])
+        W = np.array(
+            [
+                [0, 1, 0, 1],
+                [1, 0, 1, 0],
+                [0, 1, 0, 1],
+                [1, 0, 1, 0],
+            ],
+        )
 
         # centered parameterization
         with pm.Model():
-            sigma = pm.Exponential('sigma', 1)
-            phi = pm.ICAR('phi', W=W, sigma=sigma)
+            sigma = pm.Exponential("sigma", 1)
+            phi = pm.ICAR("phi", W=W, sigma=sigma)
             mu = phi
 
         # non-centered parameterization
         with pm.Model():
-            sigma = pm.Exponential('sigma', 1)
-            phi = pm.ICAR('phi', W=W)
+            sigma = pm.Exponential("sigma", 1)
+            phi = pm.ICAR("phi", W=W)
             mu = sigma * phi
 
     References

pymc/distributions/simulator.py

@@ -122,6 +122,7 @@ class Simulator(Distribution):
         def simulator_fn(rng, loc, scale, size):
             return rng.normal(loc, scale, size=size)
 
+
         with pm.Model() as m:
             loc = pm.Normal("loc", 0, 1)
             scale = pm.HalfNormal("scale", 1)