pre-commit hook for docstring format

.pre-commit-config.yaml

@@ -36,3 +36,13 @@ repos:
         stages: [pre-commit, commit-msg]
         args:
           [--ignore-words-list, "Teh,aas", --check-filenames, --skip, "*.ipynb"]
+
+  - repo: https://github.com/PyCQA/docformatter
+    # using an untagged rev for forward compatibility with pre-commit 4.0
+    # see https://github.com/PyCQA/docformatter/issues/289
+    # This should be changed back to a tag when (>1.7.5) is released
+    rev: 06907d0267368b49b9180eed423fae5697c1e909
+    hooks:
+      - id: docformatter
+        additional_dependencies: [tomli]
+        args: [--in-place, --config, ./pyproject.toml]

examples/annotation.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Bayesian Models of Annotation
 ======================================
@@ -80,9 +79,7 @@ def get_data():
 
 
 def multinomial(annotations):
-    """
-    This model corresponds to the plate diagram in Figure 1 of reference [1].
-    """
+    """This model corresponds to the plate diagram in Figure 1 of reference [1]."""
     num_classes = int(np.max(annotations)) + 1
     num_items, num_positions = annotations.shape
 
@@ -99,9 +96,7 @@ def multinomial(annotations):
 
 
 def dawid_skene(positions, annotations):
-    """
-    This model corresponds to the plate diagram in Figure 2 of reference [1].
-    """
+    """This model corresponds to the plate diagram in Figure 2 of reference [1]."""
     num_annotators = int(np.max(positions)) + 1
     num_classes = int(np.max(annotations)) + 1
     num_items, num_positions = annotations.shape
@@ -124,9 +119,7 @@ def dawid_skene(positions, annotations):
 
 
 def mace(positions, annotations):
-    """
-    This model corresponds to the plate diagram in Figure 3 of reference [1].
-    """
+    """This model corresponds to the plate diagram in Figure 3 of reference [1]."""
     num_annotators = int(np.max(positions)) + 1
     num_classes = int(np.max(annotations)) + 1
     num_items, num_positions = annotations.shape
@@ -155,9 +148,7 @@ def mace(positions, annotations):
 
 
 def hierarchical_dawid_skene(positions, annotations):
-    """
-    This model corresponds to the plate diagram in Figure 4 of reference [1].
-    """
+    """This model corresponds to the plate diagram in Figure 4 of reference [1]."""
     num_annotators = int(np.max(positions)) + 1
     num_classes = int(np.max(annotations)) + 1
     num_items, num_positions = annotations.shape
@@ -192,9 +183,7 @@ def hierarchical_dawid_skene(positions, annotations):
 
 
 def item_difficulty(annotations):
-    """
-    This model corresponds to the plate diagram in Figure 5 of reference [1].
-    """
+    """This model corresponds to the plate diagram in Figure 5 of reference [1]."""
     num_classes = int(np.max(annotations)) + 1
     num_items, num_positions = annotations.shape
 
@@ -220,9 +209,7 @@ def item_difficulty(annotations):
 
 
 def logistic_random_effects(positions, annotations):
-    """
-    This model corresponds to the plate diagram in Figure 5 of reference [1].
-    """
+    """This model corresponds to the plate diagram in Figure 5 of reference [1]."""
     num_annotators = int(np.max(positions)) + 1
     num_classes = int(np.max(annotations)) + 1
     num_items, num_positions = annotations.shape

examples/ar2.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 r"""
 Example: AR2 process
 ====================

examples/baseball.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Baseball Batting Average
 =================================
@@ -66,9 +65,8 @@
 
 
 def fully_pooled(at_bats, hits=None):
-    r"""
-    Number of hits in $K$ at bats for each player has a Binomial
-    distribution with a common probability of success, $\phi$.
+    r"""Number of hits in $K$ at bats for each player has a Binomial distribution
+    with a common probability of success, $\phi$.
 
     :param (jnp.ndarray) at_bats: Number of at bats for each player.
     :param (jnp.ndarray) hits: Number of hits for the given at bats.
@@ -82,9 +80,8 @@ def fully_pooled(at_bats, hits=None):
 
 
 def not_pooled(at_bats, hits=None):
-    r"""
-    Number of hits in $K$ at bats for each player has a Binomial
-    distribution with independent probability of success, $\phi_i$.
+    r"""Number of hits in $K$ at bats for each player has a Binomial distribution
+    with independent probability of success, $\phi_i$.
 
     :param (jnp.ndarray) at_bats: Number of at bats for each player.
     :param (jnp.ndarray) hits: Number of hits for the given at bats.
@@ -98,9 +95,9 @@ def not_pooled(at_bats, hits=None):
 
 
 def partially_pooled(at_bats, hits=None):
-    r"""
-    Number of hits has a Binomial distribution with independent
-    probability of success, $\phi_i$. Each $\phi_i$ follows a Beta
+    r"""Number of hits has a Binomial distribution with independent probability of
+    success, $\phi_i$. Each $\phi_i$ follows a Beta.
+
     distribution with concentration parameters $c_1$ and $c_2$, where
     $c_1 = m * kappa$, $c_2 = (1 - m) * kappa$, $m ~ Uniform(0, 1)$,
     and $kappa ~ Pareto(1, 1.5)$.
@@ -119,10 +116,9 @@ def partially_pooled(at_bats, hits=None):
 
 
 def partially_pooled_with_logit(at_bats, hits=None):
-    r"""
-    Number of hits has a Binomial distribution with a logit link function.
-    The logits $\alpha$ for each player is normally distributed with the
-    mean and scale parameters sharing a common prior.
+    r"""Number of hits has a Binomial distribution with a logit link function. The
+    logits $\alpha$ for each player is normally distributed with the mean and scale
+    parameters sharing a common prior.
 
     :param (jnp.ndarray) at_bats: Number of at bats for each player.
     :param (jnp.ndarray) hits: Number of hits for the given at bats.

examples/bnn.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Bayesian Neural Network
 ================================

examples/capture_recapture.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: CJS Capture-Recapture Model for Ecological Data
 ========================================================

examples/covtype.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: MCMC Methods for Tall Data
 ===================================

examples/cvae.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Conditional Variational Autoencoder in Flax
 ====================================================

examples/dais_demo.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: AutoDAIS
 =================

examples/funnel.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Neal's Funnel
 ======================

examples/gaussian_shells.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Nested Sampling for Gaussian Shells
 ============================================

examples/gp.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Gaussian Process
 =========================

examples/hmcecs.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Hamiltonian Monte Carlo with Energy Conserving Subsampling
 ===================================================================

examples/hmm.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Hidden Markov Model
 ============================

examples/hmm_enum.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Enumerate Hidden Markov Model
 ======================================

examples/holt_winters.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 r"""
 Example: Holt-Winters Exponential Smoothing
 ===========================================

examples/horseshoe_regression.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Horseshoe Regression
 =============================

examples/hsgp.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Hilbert space approximation for Gaussian processes.
 ============================================================
@@ -66,9 +65,7 @@
 
 # --- Data processing functions
 def get_labour_days(dates):
-    """
-    First monday of September
-    """
+    """First monday of September."""
     is_september = dates.dt.month.eq(9)
     is_monday = dates.dt.weekday.eq(0)
     is_first_week = dates.dt.day.le(7)
@@ -80,9 +77,7 @@ def get_labour_days(dates):
 
 
 def get_memorial_days(dates):
-    """
-    Last monday of May
-    """
+    """Last monday of May."""
     is_may = dates.dt.month.eq(5)
     is_monday = dates.dt.weekday.eq(0)
     is_last_week = dates.dt.day.ge(25)
@@ -94,9 +89,7 @@ def get_memorial_days(dates):
 
 
 def get_thanksgiving_days(dates):
-    """
-    Third thursday of November
-    """
+    """Third thursday of November."""
     is_november = dates.dt.month.eq(11)
     is_thursday = dates.dt.weekday.eq(3)
     is_third_week = dates.dt.day.between(22, 28)
@@ -172,10 +165,10 @@ def diag_spectral_density(alpha, length, L, M):
 
 
 def eigenfunctions(x, L, M):
-    """
-    The first `M` eigenfunctions of the laplacian operator in `[-L, L]`
-    evaluated at `x`. These are used for the approximation of the
-    squared exponential kernel.
+    """The first `M` eigenfunctions of the laplacian operator in `[-L, L]` evaluated
+    at `x`.
+
+    These are used for the approximation of the squared exponential kernel.
     """
     m1 = (jnp.pi / (2 * L)) * jnp.tile(L + x[:, None], M)
     m2 = jnp.diag(jnp.linspace(1, M, num=M))
@@ -190,9 +183,9 @@ def modified_bessel_first_kind(v, z):
 
 
 def diag_spectral_density_periodic(alpha, length, M):
-    """
-    Not actually a spectral density but these are used in the same
-    way. These are simply the first `M` coefficients of the low rank
+    """Not actually a spectral density but these are used in the same way.
+
+    These are simply the first `M` coefficients of the low rank
     approximation for the periodic kernel.
     """
     a = length ** (-2)
@@ -203,9 +196,7 @@ def diag_spectral_density_periodic(alpha, length, M):
 
 
 def eigenfunctions_periodic(x, w0, M):
-    """
-    Basis functions for the approximation of the periodic kernel.
-    """
+    """Basis functions for the approximation of the periodic kernel."""
     m1 = jnp.tile(w0 * x[:, None], M)
     m2 = jnp.diag(jnp.arange(M, dtype=jnp.float32))
     mw0x = m1 @ m2
@@ -216,10 +207,8 @@ def eigenfunctions_periodic(x, w0, M):
 
 # --- Approximate Gaussian processes --- #
 def approx_se_ncp(x, alpha, length, L, M):
-    """
-    Hilbert space approximation for the squared
-    exponential kernel in the non-centered parametrisation.
-    """
+    """Hilbert space approximation for the squared exponential kernel in the non-
+    centered parametrisation."""
     phi = eigenfunctions(x, L, M)
     spd = jnp.sqrt(diag_spectral_density(alpha, length, L, M))
     with plate("basis", M):
@@ -230,10 +219,8 @@ def approx_se_ncp(x, alpha, length, L, M):
 
 
 def approx_periodic_gp_ncp(x, alpha, length, w0, M):
-    """
-    Low rank approximation for the periodic squared
-    exponential kernel in the non-centered parametrisation.
-    """
+    """Low rank approximation for the periodic squared exponential kernel in the
+    non-centered parametrisation."""
     q2 = diag_spectral_density_periodic(alpha, length, M)
     cosines, sines = eigenfunctions_periodic(x, w0, M)
 

examples/mortality.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Modelling mortality over space and time
 ================================================
@@ -96,9 +95,7 @@
 
 
 def create_lookup(s1, s2):
-    """
-    Create a map between s1 indices and unique s2 indices
-    """
+    """Create a map between s1 indices and unique s2 indices."""
     lookup = np.column_stack([s1, s2])
     lookup = np.unique(lookup, axis=0)
     lookup = lookup[lookup[:, 1].argsort()]

examples/neutra.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Neural Transport
 =========================

examples/ode.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: Predator-Prey Model
 ============================
@@ -40,8 +39,9 @@
 
 
 def dz_dt(z, t, theta):
-    """
-    Lotka–Volterra equations. Real positive parameters `alpha`, `beta`, `gamma`, `delta`
+    """Lotka–Volterra equations.
+
+    Real positive parameters `alpha`, `beta`, `gamma`, `delta`
     describes the interaction of two species.
     """
     u = z[0]

examples/prodlda.py

@@ -1,6 +1,5 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
-
 """
 Example: ProdLDA with Flax and Haiku
 ====================================