Write docstrings in imperative mood (D401)

Author: Armavica

pymc/backends/__init__.py

@@ -101,7 +101,7 @@ def _init_trace(
     model: Model,
     trace_vars: list[TensorVariable] | None = None,
 ) -> BaseTrace:
-    """Initializes a trace backend for a chain."""
+    """Initialize a trace backend for a chain."""
     strace: BaseTrace
     if trace is None:
         strace = NDArray(model=model, vars=trace_vars)
@@ -126,7 +126,7 @@ def init_traces(
     model: Model,
     trace_vars: list[TensorVariable] | None = None,
 ) -> tuple[RunType | None, Sequence[IBaseTrace]]:
-    """Initializes a trace recorder for each chain."""
+    """Initialize a trace recorder for each chain."""
     if HAS_MCB and isinstance(backend, Backend):
         return init_chain_adapters(
             backend=backend,

pymc/backends/mcbackend.py

@@ -43,7 +43,7 @@
 
 
 def find_data(pmodel: Model) -> list[mcb.DataVariable]:
-    """Extracts data variables from a model."""
+    """Extract data variables from a model."""
     observed_rvs = {pmodel.rvs_to_values[rv] for rv in pmodel.observed_RVs}
     dvars = []
     # All data containers are named vars!
@@ -131,7 +131,7 @@ def get_values(self, varname: str, burn=0, thin=1) -> np.ndarray:
         return self._chain.get_draws(varname, slice(burn, None, thin))
 
     def _get_stats(self, fname: str, slc: slice) -> np.ndarray:
-        """Wraps `self._chain.get_stats` but unpickles automatically."""
+        """Wrap `self._chain.get_stats` but unpickle automatically."""
         values = self._chain.get_stats(fname, slc)
         # Unpickle object stats
         if fname in self._statsbj.object_stats:

pymc/data.py

@@ -53,7 +53,7 @@
 
 
 def get_data(filename):
-    """Returns a BytesIO object for a package data file.
+    """Return a BytesIO object for a package data file.
 
     Parameters
     ----------
@@ -224,7 +224,7 @@ def determine_coords(
     dims: Sequence[str | None] | None = None,
     coords: dict[str, Sequence | np.ndarray] | None = None,
 ) -> tuple[dict[str, Sequence | np.ndarray], Sequence[str | None]]:
-    """Determines coordinate values from data or the model (via ``dims``)."""
+    """Determine coordinate values from data or the model (via ``dims``)."""
     if coords is None:
         coords = {}
 
@@ -343,7 +343,7 @@ def Data(
     mutable: bool | None = None,
     **kwargs,
 ) -> SharedVariable | TensorConstant:
-    """Data container that registers a data variable with the model.
+    """Create a data container that registers a data variable with the model.
 
     Depending on the ``mutable`` setting (default: True), the variable
     is registered as a :class:`~pytensor.compile.sharedvalue.SharedVariable`,

pymc/distributions/dist_math.py

@@ -90,7 +90,7 @@ def check_icdf_value(expr: Variable, value: Variable) -> Variable:
 
 
 def logpow(x, m):
-    """Calculates log(x**m) since m*log(x) will fail when m, x = 0."""
+    """Calculate log(x**m) since m*log(x) will fail when m, x = 0."""
     # return m * log(x)
     return pt.switch(pt.eq(x, 0), pt.switch(pt.eq(m, 0), 0.0, -np.inf), m * pt.log(x))
 
@@ -108,7 +108,7 @@ def betaln(x, y):
 
 
 def std_cdf(x):
-    """Calculates the standard normal cumulative distribution function."""
+    """Calculate the standard normal cumulative distribution function."""
     return 0.5 + 0.5 * pt.erf(x / pt.sqrt(2.0))
 
 
@@ -400,7 +400,7 @@ def multigammaln(a, p):
 
 
 def log_i0(x):
-    """Calculates the logarithm of the 0 order modified Bessel function of the first kind."""
+    """Calculate the logarithm of the 0 order modified Bessel function of the first kind."""
     return pt.switch(
         pt.lt(x, 5),
         pt.log1p(

pymc/distributions/distribution.py

@@ -392,7 +392,7 @@ def update(self, node: Apply) -> dict[Variable, Variable]:
         return collect_default_updates_inner_fgraph(node)
 
     def batch_ndim(self, node: Apply) -> int:
-        """Number of dimensions of the distribution's batch shape."""
+        """Return the number of dimensions of the distribution's batch shape."""
         out_ndim = max(getattr(out.type, "ndim", 0) for out in node.outputs)
         return out_ndim - self.ndim_supp
 
@@ -438,7 +438,7 @@ def __new__(
         default_transform=UNSET,
         **kwargs,
     ) -> TensorVariable:
-        """Adds a tensor variable corresponding to a PyMC distribution to the current model.
+        """Add a tensor variable corresponding to a PyMC distribution to the current model.
 
         Note that all remaining kwargs must be compatible with ``.dist()``
 
@@ -531,7 +531,7 @@ def dist(
         shape: Shape | None = None,
         **kwargs,
     ) -> TensorVariable:
-        """Creates a tensor variable corresponding to the `cls` distribution.
+        """Create a tensor variable corresponding to the `cls` distribution.
 
         Parameters
         ----------
@@ -604,8 +604,7 @@ def _support_point(op, rv, *rv_inputs) -> TensorVariable:
 
 
 def support_point(rv: TensorVariable) -> TensorVariable:
-    """Method for choosing a representative point/value
-    that can be used to start optimization or MCMC sampling.
+    """Choose a representative point/value that can be used to start optimization or MCMC sampling.
 
     The only parameter to this function is the RandomVariable
     for which the value is to be derived.

pymc/distributions/mixture.py

@@ -555,9 +555,9 @@ def dist(cls, w, mu, sigma=None, tau=None, **kwargs):
 
 
 def _zero_inflated_mixture(*, name, nonzero_p, nonzero_dist, **kwargs):
-    """Helper function to create a zero-inflated mixture.
+    """Create a zero-inflated mixture (helper function).
 
-    If name is `None`, this function returns an unregistered variable
+    If name is `None`, this function returns an unregistered variable.
     """
     nonzero_p = pt.as_tensor_variable(nonzero_p)
     weights = pt.stack([1 - nonzero_p, nonzero_p], axis=-1)
@@ -799,7 +799,7 @@ def dist(cls, psi, mu=None, alpha=None, p=None, n=None, **kwargs):
 
 
 def _hurdle_mixture(*, name, nonzero_p, nonzero_dist, dtype, max_n_steps=10_000, **kwargs):
-    """Helper function to create a hurdle mixtures.
+    """Create a hurdle mixtures (helper function).
 
     If name is `None`, this function returns an unregistered variable
 

pymc/distributions/multivariate.py

@@ -342,7 +342,7 @@ def precision_mv_normal_logp(op: PrecisionMvNormalRV, value, rng, size, mean, ta
 
 @node_rewriter(tracks=[MvNormalRV])
 def mv_normal_to_precision_mv_normal(fgraph, node):
-    """Replaces MvNormal(mu, inv(tau)) -> PrecisionMvNormal(mu, tau).
+    """Replace MvNormal(mu, inv(tau)) -> PrecisionMvNormal(mu, tau).
 
     This is introduced in logprob rewrites to provide a more efficient logp for a MvNormal
     that is defined by a precision matrix.
@@ -2153,7 +2153,8 @@ def make_node(self, rng, size, mu, W, alpha, tau, W_is_valid):
 
     @classmethod
     def rng_fn(cls, rng: np.random.RandomState, mu, W, alpha, tau, W_is_valid, size):
-        """
+        """Sample a numeric random variate.
+
         Implementation of algorithm from paper
         Havard Rue, 2001. "Fast sampling of Gaussian Markov random fields,"
         Journal of the Royal Statistical Society Series B, Royal Statistical Society,

pymc/distributions/shape_utils.py

@@ -149,7 +149,7 @@ def convert_size(size: Size) -> StrongSize | None:
 
 
 def shape_from_dims(dims: StrongDims, model) -> StrongShape:
-    """Determines shape from a `dims` tuple.
+    """Determine shape from a `dims` tuple.
 
     Parameters
     ----------
@@ -178,7 +178,7 @@ def find_size(
     size: StrongSize | None,
     ndim_supp: int,
 ) -> StrongSize | None:
-    """Determines the size keyword argument for creating a Distribution.
+    """Determine the size keyword argument for creating a Distribution.
 
     Parameters
     ----------
@@ -436,7 +436,11 @@ def get_support_shape_1d(
     observed: Any | None = None,
     support_shape_offset: int = 0,
 ) -> TensorVariable | None:
-    """Helper function for cases when you just care about one dimension."""
+    """
+    Extract the support shapes from shape / dims / observed information.
+
+    Helper function for cases when you just care about one dimension.
+    """
     support_shape_tuple = get_support_shape(
         support_shape=(support_shape,) if support_shape is not None else None,
         shape=shape,

pymc/gp/cov.py

@@ -114,7 +114,7 @@ def __pow__(self, other) -> "Exponentiated":
         return Exponentiated(self, other)
 
     def __array_wrap__(self, result):
-        """Required to allow radd/rmul by numpy arrays."""
+        """Allow radd/rmul by numpy arrays."""
         result = np.squeeze(result)
         if len(result.shape) <= 1:
             result = result.reshape(1, 1)
@@ -234,7 +234,7 @@ def __init__(self, factor_list: Sequence):
                 self._factor_list.append(factor)
 
     def _merge_factors_cov(self, X, Xs=None, diag=False):
-        """Called to evaluate either all the sums or all the
+        """Evaluate either all the sums or all the
         products of kernels that are possible to evaluate.
         """
         factor_list = []
@@ -264,12 +264,12 @@ def _merge_factors_cov(self, X, Xs=None, diag=False):
         return factor_list
 
     def _merge_factors_psd(self, omega):
-        """Called to evaluatate spectral densities of combination kernels when possible.
+        """Evaluate spectral densities of combination kernels when possible.
 
-        Implements
-        a more restricted set of rules than `_merge_factors_cov` -- just additivity of stationary
-        covariances with defined power spectral densities and multiplication by scalars.  Also, the
-        active_dims for all covariances in the sum must be the same.
+        Implements a more restricted set of rules than `_merge_factors_cov` --
+        just additivity of stationary covariances with defined power spectral
+        densities and multiplication by scalars.  Also, the active_dims for all
+        covariances in the sum must be the same.
         """
         factor_list = []
         for factor in self._factor_list:

pymc/gp/gp.py

@@ -47,7 +47,7 @@
 
 
 def _handle_sigma_noise_parameters(sigma, noise):
-    """Helper function for transition of 'noise' parameter to be named 'sigma'."""
+    """Help transition of 'noise' parameter to be named 'sigma'."""
     if (sigma is None and noise is None) or (sigma is not None and noise is not None):
         raise ValueError("'sigma' argument must be specified.")
 
@@ -177,7 +177,7 @@ def _build_prior(
 
     def prior(self, name, X, n_outputs=1, reparameterize=True, jitter=JITTER_DEFAULT, **kwargs):
         R"""
-        Returns the GP prior distribution evaluated over the input
+        Return the GP prior distribution evaluated over the input
         locations `X`.
 
         This is the prior probability over the space
@@ -250,7 +250,7 @@ def _build_conditional(self, Xnew, X, f, cov_total, mean_total, jitter):
 
     def conditional(self, name, Xnew, given=None, jitter=JITTER_DEFAULT, **kwargs):
         R"""
-        Returns the conditional distribution evaluated over new input
+        Return the conditional distribution evaluated over new input
         locations `Xnew`.
 
         Given a set of function values `f` that
@@ -350,7 +350,7 @@ def _build_prior(self, name, X, reparameterize=True, jitter=JITTER_DEFAULT, **kw
 
     def prior(self, name, X, reparameterize=True, jitter=JITTER_DEFAULT, **kwargs):
         R"""
-        Returns the TP prior distribution evaluated over the input
+        Return the TP prior distribution evaluated over the input
         locations `X`.
 
         This is the prior probability over the space
@@ -394,7 +394,7 @@ def _build_conditional(self, Xnew, X, f, jitter):
 
     def conditional(self, name, Xnew, jitter=JITTER_DEFAULT, **kwargs):
         R"""
-        Returns the conditional distribution evaluated over new input
+        Return the conditional distribution evaluated over new input
         locations `Xnew`.
 
         Given a set of function values `f` that
@@ -487,7 +487,7 @@ def marginal_likelihood(
         **kwargs,
     ):
         R"""
-        Returns the marginal likelihood distribution, given the input
+        Return the marginal likelihood distribution, given the input
         locations `X` and the data `y`.
 
         This is the integral over the product of the GP prior and a normal likelihood.
@@ -594,7 +594,7 @@ def conditional(
         self, name, Xnew, pred_noise=False, given=None, jitter=JITTER_DEFAULT, **kwargs
     ):
         R"""
-        Returns the conditional distribution evaluated over new input
+        Return the conditional distribution evaluated over new input
         locations `Xnew`.
 
         Given a set of function values `f` that the GP prior was over, the
@@ -814,7 +814,7 @@ def marginal_likelihood(
         self, name, X, Xu, y, sigma=None, noise=None, jitter=JITTER_DEFAULT, **kwargs
     ):
         R"""
-        Returns the approximate marginal likelihood distribution, given the input
+        Return the approximate marginal likelihood distribution, given the input
         locations `X`, inducing point locations `Xu`, data `y`, and white noise
         standard deviations `sigma`.
 
@@ -906,7 +906,7 @@ def conditional(
         self, name, Xnew, pred_noise=False, given=None, jitter=JITTER_DEFAULT, **kwargs
     ):
         R"""
-        Returns the approximate conditional distribution of the GP evaluated over
+        Return the approximate conditional distribution of the GP evaluated over
         new input locations `Xnew`.
 
         Parameters
@@ -1019,7 +1019,7 @@ def _build_prior(self, name, Xs, jitter, **kwargs):
 
     def prior(self, name, Xs, jitter=JITTER_DEFAULT, **kwargs):
         """
-        Returns the prior distribution evaluated over the input
+        Return the prior distribution evaluated over the input
         locations `Xs`.
 
         Parameters
@@ -1065,7 +1065,7 @@ def _build_conditional(self, Xnew, jitter):
 
     def conditional(self, name, Xnew, jitter=JITTER_DEFAULT, **kwargs):
         """
-        Returns the conditional distribution evaluated over new input
+        Return the conditional distribution evaluated over new input
         locations `Xnew`.
 
         `Xnew` will be split by columns and fed to the relevant
@@ -1188,7 +1188,7 @@ def _check_inputs(self, Xs, y):
 
     def marginal_likelihood(self, name, Xs, y, sigma, is_observed=True, **kwargs):
         """
-        Returns the marginal likelihood distribution, given the input
+        Return the marginal likelihood distribution, given the input
         locations `cartesian(*Xs)` and the data `y`.
 
         Parameters
@@ -1267,7 +1267,7 @@ def _build_conditional(self, Xnew, diag, pred_noise):
 
     def conditional(self, name, Xnew, pred_noise=False, diag=False, **kwargs):
         """
-        Returns the conditional distribution evaluated over new input
+        Return the conditional distribution evaluated over new input
         locations `Xnew`, just as in `Marginal`.
 
         `Xnew` will be split by columns and fed to the relevant

pymc/gp/hsgp_approx.py

@@ -93,8 +93,7 @@ def calc_basis_periodic(
 def approx_hsgp_hyperparams(
     x_range: list[float], lengthscale_range: list[float], cov_func: str
 ) -> tuple[int, float]:
-    """Utility function that uses heuristics to recommend minimum `m` and `c` values,
-    based on recommendations from Ruitort-Mayol et. al.
+    """Use heuristics to recommend minimum `m` and `c` values, based on recommendations from Ruitort-Mayol et. al.
 
     In practice, you need to choose `c` large enough to handle the largest lengthscales,
     and `m` large enough to accommodate the smallest lengthscales.  Use your prior on the
@@ -427,7 +426,7 @@ def prior(
         **kwargs,
     ):  # type: ignore
         R"""
-        Returns the (approximate) GP prior distribution evaluated over the input locations `X`.
+        Return the (approximate) GP prior distribution evaluated over the input locations `X`.
         For usage examples, refer to `pm.gp.Latent`.
 
         Parameters
@@ -490,7 +489,7 @@ def _build_conditional(self, Xnew):
 
     def conditional(self, name: str, Xnew: TensorLike, dims: str | None = None):  # type: ignore
         R"""
-        Returns the (approximate) conditional distribution evaluated over new input locations
+        Return the (approximate) conditional distribution evaluated over new input locations
         `Xnew`.
 
         Parameters
@@ -685,7 +684,7 @@ def prior_linearized(self, X: TensorLike):
 
     def prior(self, name: str, X: TensorLike, dims: str | None = None):  # type: ignore
         R"""
-        Returns the (approximate) GP prior distribution evaluated over the input locations `X`.
+        Return the (approximate) GP prior distribution evaluated over the input locations `X`.
         For usage examples, refer to `pm.gp.Latent`.
 
         Parameters
@@ -736,7 +735,7 @@ def _build_conditional(self, Xnew):
 
     def conditional(self, name: str, Xnew: TensorLike, dims: str | None = None):  # type: ignore
         R"""
-        Returns the (approximate) conditional distribution evaluated over new input locations
+        Return the (approximate) conditional distribution evaluated over new input locations
         `Xnew`.
 
         Parameters