Run formatting

src/jnotype/__init__.py

@@ -1,4 +1,5 @@
 """Exploratory analysis of binary data."""
+
 import jnotype.bmm as bmm
 import jnotype.datasets as datasets
 import jnotype.sampling as sampling

src/jnotype/_csp.py

@@ -1,4 +1,5 @@
 """Cumulative shrinkage prior."""
+
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp

src/jnotype/_factor_analysis/_gibbs_backend.py

@@ -1,5 +1,6 @@
 """Sampling steps for all variables, apart from variances
 attributed to latent traits, which are sampled with CSP module."""
+
 from typing import Callable
 
 import jax

src/jnotype/_factor_analysis/_inference.py

@@ -1,4 +1,5 @@
 """Sampling from the posterior distribution."""
+
 from typing import Callable
 
 from jaxtyping import Float, Array

src/jnotype/_factor_analysis/_simulate.py

@@ -1,4 +1,5 @@
 """Simulate data sets."""
+
 from jaxtyping import Float, Array
 
 import jax.numpy as jnp

src/jnotype/_utils.py

@@ -3,6 +3,7 @@
 This file should be as small as possible.
 Appearing themes should be refactored and placed
 into separate modules."""
+
 import jax
 
 

src/jnotype/_variance.py

@@ -1,4 +1,5 @@
 """Utilities for sampling variances."""
+
 from jax import random
 import jax
 import jax.numpy as jnp

src/jnotype/bmm/__init__.py

@@ -1,4 +1,5 @@
 """Bernoulli Mixture Model."""
+
 from jnotype.bmm._em import expectation_maximization
 from jnotype.bmm._gibbs import (
     sample_mixing,

src/jnotype/bmm/_em.py

@@ -1,4 +1,5 @@
 """The Expectation-Maximization algorithm for Bernoulli Mixture Model."""
+
 import dataclasses
 import time
 from typing import Optional
@@ -68,7 +69,11 @@ def em_step(
     observed: Int[Array, "N K"],
     mixing: Float[Array, "K B"],
     proportions: Float[Array, " B"],
-) -> tuple[Float[Array, "N B"], Float[Array, "K B"], Float[Array, " B"],]:
+) -> tuple[
+    Float[Array, "N B"],
+    Float[Array, "K B"],
+    Float[Array, " B"],
+]:
     """The E and M step combined, for better JIT compiler optimisation.
 
     Args:

src/jnotype/bmm/_gibbs.py

@@ -1,4 +1,5 @@
 """Sampling cluster labels and proportions."""
+
 from typing import Optional, Sequence
 
 import jax

src/jnotype/checks/_histograms.py

@@ -1,4 +1,5 @@
 """Plotting histograms of data."""
+
 from typing import Sequence, Union, Literal
 
 import matplotlib.pyplot as plt

src/jnotype/datasets/__init__.py

@@ -1,4 +1,5 @@
 """Data sets."""
+
 from jnotype.datasets._simulation import BlockImagesSampler
 
 __all__ = [

src/jnotype/datasets/_simulation/__init__.py

@@ -1,4 +1,5 @@
 """Simulated data sets."""
+
 from jnotype.datasets._simulation._block_images import BlockImagesSampler
 
 __all__ = [

src/jnotype/datasets/_simulation/_block_images.py

@@ -1,4 +1,5 @@
 """Simulation of binary images using Bernoulli mixture model."""
+
 from typing import Optional
 
 from jaxtyping import Array, Float, Int

src/jnotype/logistic/_binary_latent.py

@@ -1,4 +1,5 @@
 """Sample binary latent variables."""
+
 from functools import partial
 
 import jax

src/jnotype/logistic/_polyagamma.py

@@ -1,4 +1,5 @@
 """Logistic regression sampling utilities using Pólya-Gamma augmentation."""
+
 from jax import random
 import jax
 import jax.numpy as jnp
@@ -53,9 +54,9 @@ def _sample_coefficients(
     precision_matrices: Float[Array, "features covariates covariates"] = jax.vmap(
         jnp.diag
     )(jnp.reciprocal(prior_variance))
-    posterior_covariances: Float[
-        Array, "features covariates covariates"
-    ] = jnp.linalg.inv(x_omega_x + precision_matrices)
+    posterior_covariances: Float[Array, "features covariates covariates"] = (
+        jnp.linalg.inv(x_omega_x + precision_matrices)
+    )
 
     kappa: Float[Array, "points features"] = jnp.asarray(observed, dtype=float) - 0.5
 

src/jnotype/logistic/_structure.py

@@ -1,4 +1,5 @@
 """Sample structure (spike/slab distinction) variables."""
+
 from typing import Union
 
 import jax

src/jnotype/logistic/logreg.py

@@ -1,4 +1,5 @@
 """Logistic regression utilities."""
+
 import jax
 import jax.numpy as jnp
 from jaxtyping import Int, Float, Array

src/jnotype/pyramids/_sampler_csp.py

@@ -1,6 +1,7 @@
 """Sampler for two-layer Bayesian pyramids
 with cumulative shrinkage process (CSP) prior
 on latent binary codes."""
+
 from typing import Optional, Sequence, Union, NewType
 
 import jax

src/jnotype/pyramids/_sampler_fixed.py

@@ -1,5 +1,6 @@
 """Sampler for two-layer Bayesian pyramids
 with fixed number of latent binary codes."""
+
 from typing import Optional, Sequence, Union, NewType
 
 import jax

src/jnotype/sampling/__init__.py

@@ -1,4 +1,5 @@
 """Generic utilities for sampling."""
+
 from jnotype.sampling._chunker import (
     DatasetInterface,
     ListDataset,

src/jnotype/sampling/_chunker.py

@@ -1,4 +1,5 @@
 """Utilities for saving samples in chunks, to limit RAM usage."""
+
 import abc
 
 from datetime import datetime

src/jnotype/sampling/_sampler.py

@@ -1,4 +1,5 @@
 """Generic Gibbs sampler."""
+
 import abc
 import logging
 import time