run ruff linter + formatter

Author: LarsKue

.gitignore

@@ -32,4 +32,4 @@ docs/
 .tox
 
 # MacOS
-.DS_Store
+.DS_Store

CITATION.cff

@@ -69,4 +69,4 @@ preferred-citation:
   type: article
   url: "https://joss.theoj.org/papers/10.21105/joss.05702"
   volume: 8
-title: "BayesFlow: Amortized Bayesian Workflows With Neural Networks"
+title: "BayesFlow: Amortized Bayesian Workflows With Neural Networks"

bayesflow/__init__.py

@@ -1,4 +1,3 @@
-
 from . import (
     approximators,
     configurators,

bayesflow/approximators/__init__.py

@@ -1,2 +1 @@
-
 from .approximator import Approximator

bayesflow/approximators/approximator.py

@@ -1,4 +1,3 @@
-
 import keras
 from keras.saving import register_keras_serializable
 
@@ -20,7 +19,7 @@
 @register_keras_serializable(package="bayesflow.amortizers")
 class Approximator(BaseApproximator):
     def __init__(self, **kwargs):
-        """ The main workhorse for learning amortized neural approximators for distributions arising
+        """The main workhorse for learning amortized neural approximators for distributions arising
         in inverse problems and Bayesian inference (e.g., posterior distributions, likelihoods, marginal
         likelihoods).
 
@@ -64,14 +63,16 @@ def __init__(self, **kwargs):
         if "configurator" not in kwargs:
             # try to set up a default configurator
             if "inference_variables" not in kwargs:
-                raise ValueError(f"You must specify either a configurator or arguments for the default configurator.")
+                raise ValueError("You must specify either a configurator or arguments for the default configurator.")
 
             inference_variables = kwargs.pop("inference_variables")
             inference_conditions = kwargs.pop("inference_conditions", None)
             summary_variables = kwargs.pop("summary_variables", None)
             summary_conditions = kwargs.pop("summary_conditions", None)
 
-            kwargs["configurator"] = Configurator(inference_variables, inference_conditions, summary_variables, summary_conditions)
+            kwargs["configurator"] = Configurator(
+                inference_variables, inference_conditions, summary_variables, summary_conditions
+            )
 
         kwargs.setdefault("summary_network", None)
         super().__init__(**kwargs)

bayesflow/approximators/base_approximator.py

@@ -1,4 +1,3 @@
-
 import keras
 from keras.saving import (
     deserialize_keras_object,
@@ -15,15 +14,23 @@
 
 @register_keras_serializable(package="bayesflow.approximators")
 class BaseApproximator(keras.Model):
-    def __init__(self, inference_network: InferenceNetwork, summary_network: SummaryNetwork, configurator: BaseConfigurator, **kwargs):
+    def __init__(
+        self,
+        inference_network: InferenceNetwork,
+        summary_network: SummaryNetwork,
+        configurator: BaseConfigurator,
+        **kwargs,
+    ):
         super().__init__(**keras_kwargs(kwargs))
         self.inference_network = inference_network
         self.summary_network = summary_network
         self.configurator = configurator
 
     @classmethod
     def from_config(cls, config: dict, custom_objects=None) -> "BaseApproximator":
-        config["inference_network"] = deserialize_keras_object(config["inference_network"], custom_objects=custom_objects)
+        config["inference_network"] = deserialize_keras_object(
+            config["inference_network"], custom_objects=custom_objects
+        )
         config["summary_network"] = deserialize_keras_object(config["summary_network"], custom_objects=custom_objects)
         config["configurator"] = deserialize_keras_object(config["configurator"], custom_objects=custom_objects)
 
@@ -63,10 +70,12 @@ def evaluate(self, *args, **kwargs):
 
         if val_logs is None:
             # https://github.com/keras-team/keras/issues/19835
-            warnings.warn(f"Found no validation logs due to a bug in keras. "
-                          f"Applying workaround, but incorrect loss values may be logged. "
-                          f"If possible, increase the size of your dataset, "
-                          f"or lower the number of validation steps used.")
+            warnings.warn(
+                "Found no validation logs due to a bug in keras. "
+                "Applying workaround, but incorrect loss values may be logged. "
+                "If possible, increase the size of your dataset, "
+                "or lower the number of validation steps used."
+            )
 
             val_logs = {}
 
@@ -103,16 +112,18 @@ def compute_metrics(self, data: dict[str, Tensor], stage: str = "training") -> d
         return metrics | summary_metrics | inference_metrics
 
     def compute_loss(self, *args, **kwargs):
-        raise RuntimeError(f"Use compute_metrics()['loss'] instead.")
+        raise RuntimeError("Use compute_metrics()['loss'] instead.")
 
     def fit(self, *args, **kwargs):
         if not self.built:
             try:
                 dataset = kwargs.get("x") or args[0]
                 self.build_from_data(dataset[0])
             except Exception:
-                raise RuntimeError(f"Could not automatically build the approximator. Please pass a dataset as the "
-                                   f"first argument to `approximator.fit()` or manually call `approximator.build()` "
-                                   f"with a dictionary specifying your data shapes.")
+                raise RuntimeError(
+                    "Could not automatically build the approximator. Please pass a dataset as the "
+                    "first argument to `approximator.fit()` or manually call `approximator.build()` "
+                    "with a dictionary specifying your data shapes."
+                )
 
         return super().fit(*args, **kwargs)

bayesflow/approximators/jax_approximator.py

@@ -1,4 +1,3 @@
-
 import jax
 import keras
 
@@ -13,7 +12,14 @@ def train_step(self, *args, **kwargs):
     def test_step(self, *args, **kwargs):
         return self.stateless_test_step(*args, **kwargs)
 
-    def stateless_compute_metrics(self, trainable_variables: any, non_trainable_variables: any, metrics_variables: any, data: dict[str, Tensor], stage: str = "training") -> (Tensor, tuple):
+    def stateless_compute_metrics(
+        self,
+        trainable_variables: any,
+        non_trainable_variables: any,
+        metrics_variables: any,
+        data: dict[str, Tensor],
+        stage: str = "training",
+    ) -> (Tensor, tuple):
         """
         Things we do for jax:
         1. Accept trainable variables as the first argument
@@ -47,11 +53,13 @@ def stateless_train_step(self, state: tuple, data: dict[str, Tensor]) -> (dict[s
 
         grad_fn = jax.value_and_grad(self.stateless_compute_metrics, has_aux=True)
 
-        (loss, aux), grads = grad_fn(trainable_variables, non_trainable_variables, metrics_variables, data, stage="training")
+        (loss, aux), grads = grad_fn(
+            trainable_variables, non_trainable_variables, metrics_variables, data, stage="training"
+        )
         metrics, non_trainable_variables, metrics_variables = aux
 
-        trainable_variables, optimizer_variables = (
-            self.optimizer.stateless_apply(optimizer_variables, grads, trainable_variables)
+        trainable_variables, optimizer_variables = self.optimizer.stateless_apply(
+            optimizer_variables, grads, trainable_variables
         )
 
         metrics_variables = self._update_loss(loss, metrics_variables)
@@ -62,7 +70,9 @@ def stateless_train_step(self, state: tuple, data: dict[str, Tensor]) -> (dict[s
     def stateless_test_step(self, state: tuple, data: dict[str, Tensor]) -> (dict[str, Tensor], tuple):
         trainable_variables, non_trainable_variables, metrics_variables = state
 
-        loss, aux = self.stateless_compute_metrics(trainable_variables, non_trainable_variables, metrics_variables, data, stage="validation")
+        loss, aux = self.stateless_compute_metrics(
+            trainable_variables, non_trainable_variables, metrics_variables, data, stage="validation"
+        )
         metrics, non_trainable_variables, metrics_variables = aux
 
         metrics_variables = self._update_loss(loss, metrics_variables)

bayesflow/approximators/numpy_approximator.py

@@ -1,11 +1,8 @@
-
-import numpy as np
-
 from bayesflow.types import Tensor
 
 from .base_approximator import BaseApproximator
 
 
 class NumpyApproximator(BaseApproximator):
     def train_step(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
-        raise NotImplementedError(f"Keras currently has no support for numpy training.")
+        raise NotImplementedError("Keras currently has no support for numpy training.")

bayesflow/approximators/tensorflow_approximator.py

@@ -1,4 +1,3 @@
-
 import tensorflow as tf
 
 from bayesflow.types import Tensor

bayesflow/approximators/torch_approximator.py

@@ -1,4 +1,3 @@
-
 import torch
 
 from .base_approximator import BaseApproximator

bayesflow/configurators/__init__.py

@@ -1,3 +1,2 @@
-
 from .base_configurator import BaseConfigurator
 from .configurator import Configurator

bayesflow/configurators/base_configurator.py

@@ -1,4 +1,3 @@
-
 from keras.saving import register_keras_serializable
 
 from bayesflow.types import Tensor

bayesflow/configurators/configurator.py

@@ -1,4 +1,3 @@
-
 from keras.saving import register_keras_serializable
 
 from bayesflow.types import Tensor
@@ -14,7 +13,7 @@ def __init__(
         inference_variables: list[str],
         inference_conditions: list[str] = None,
         summary_variables: list[str] = None,
-        summary_conditions: list[str] = None
+        summary_conditions: list[str] = None,
     ):
         self.inference_variables = inference_variables
         self.inference_conditions = inference_conditions or []
@@ -30,7 +29,7 @@ def get_config(self) -> dict:
             "inference_variables": self.inference_variables,
             "inference_conditions": self.inference_conditions,
             "summary_variables": self.summary_variables,
-            "summary_conditions": self.summary_conditions
+            "summary_conditions": self.summary_conditions,
         }
 
     def configure_inference_variables(self, data: dict[str, Tensor]) -> Tensor | None:

bayesflow/datasets/__init__.py

@@ -1,5 +1,3 @@
-
 from .offline_dataset import OfflineDataset
 from .online_dataset import OnlineDataset
 from .rounds_dataset import RoundsDataset
-

bayesflow/datasets/offline_dataset.py

@@ -1,4 +1,3 @@
-
 import keras
 import math
 
@@ -7,6 +6,7 @@ class OfflineDataset(keras.utils.PyDataset):
     """
     A dataset that is pre-simulated and stored in memory.
     """
+
     def __init__(self, data: dict, batch_size: int, **kwargs):
         super().__init__(**kwargs)
         self.batch_size = batch_size
@@ -17,7 +17,7 @@ def __init__(self, data: dict, batch_size: int, **kwargs):
         self.shuffle()
 
     def __getitem__(self, item: int) -> (dict, dict):
-        """ Get a batch of pre-simulated data """
+        """Get a batch of pre-simulated data"""
         item = slice(item * self.batch_size, (item + 1) * self.batch_size)
         item = self.indices[item]
 
@@ -30,5 +30,5 @@ def on_epoch_end(self) -> None:
         self.shuffle()
 
     def shuffle(self) -> None:
-        """ Shuffle the dataset in-place. """
+        """Shuffle the dataset in-place."""
         self.indices = keras.random.shuffle(self.indices)

bayesflow/datasets/online_dataset.py

@@ -1,4 +1,3 @@
-
 import keras
 
 from bayesflow.simulators.simulator import Simulator
@@ -8,12 +7,14 @@ class OnlineDataset(keras.utils.PyDataset):
     """
     A dataset that is generated on-the-fly.
     """
+
     def __init__(self, simulator: Simulator, batch_size: int, **kwargs):
         super().__init__(**kwargs)
 
         if kwargs.get("use_multiprocessing"):
             # keras workaround: https://github.com/keras-team/keras/issues/19346
             import multiprocessing as mp
+
             mp.set_start_method("spawn", force=True)
 
         self.simulator = simulator

bayesflow/datasets/rounds_dataset.py

@@ -1,4 +1,3 @@
-
 import keras
 
 from bayesflow.simulators.simulator import Simulator
@@ -8,12 +7,14 @@ class RoundsDataset(keras.utils.PyDataset):
     """
     A dataset that is generated on-the-fly at the beginning of every n-th epoch.
     """
+
     def __init__(self, simulator: Simulator, batch_size: int, batches_per_epoch: int, epochs_per_round: int, **kwargs):
         super().__init__(**kwargs)
 
         if kwargs.get("use_multiprocessing"):
             # keras workaround: https://github.com/keras-team/keras/issues/19346
             import multiprocessing as mp
+
             mp.set_start_method("spawn", force=True)
 
         self.simulator = simulator
@@ -27,7 +28,7 @@ def __init__(self, simulator: Simulator, batch_size: int, batches_per_epoch: int
         self.regenerate()
 
     def __getitem__(self, item: int) -> (dict, dict):
-        """ Get a batch of pre-simulated data """
+        """Get a batch of pre-simulated data"""
         return self.data[item]
 
     @property
@@ -41,5 +42,5 @@ def on_epoch_end(self) -> None:
             self.regenerate()
 
     def regenerate(self) -> None:
-        """ Sample new batches of data from the joint distribution unconditionally """
+        """Sample new batches of data from the joint distribution unconditionally"""
         self.data = [self.simulator.sample((self.batch_size,)) for _ in range(self.batches_per_epoch)]

bayesflow/distributions/__init__.py

@@ -1,3 +1,2 @@
-
 from .distribution import Distribution
 from .diagonal_normal import DiagonalNormal

bayesflow/distributions/diagonal_normal.py

@@ -1,4 +1,3 @@
-
 import math
 
 import keras
@@ -16,6 +15,7 @@ class DiagonalNormal(Distribution):
         - ``_log_unnormalized_prob`` method is used as a loss function
         - ``log_prob`` is used for density computation
     """
+
     def __init__(self, mean: float | Tensor = 0.0, std: float | Tensor = 1.0, **kwargs):
         super().__init__(**kwargs)
         self.mean = mean

bayesflow/distributions/distribution.py

@@ -1,4 +1,3 @@
-
 import keras
 
 from bayesflow.types import Shape, Tensor

bayesflow/networks/__init__.py

@@ -1,11 +1,9 @@
-
 from .coupling_flow import CouplingFlow
 from .deep_set import DeepSet
 from .flow_matching import FlowMatching
+from .inference_network import InferenceNetwork
 from .mlp import MLP
 from .resnet import ResNet
 from .lstnet import LSTNet
+from .summary_network import SummaryNetwork
 from .transformers import SetTransformer
-
-from .inference_network import InferenceNetwork
-from .summary_network import SummaryNetwork

bayesflow/networks/coupling_flow/__init__.py

@@ -1,2 +1 @@
-
 from .coupling_flow import CouplingFlow

bayesflow/networks/coupling_flow/actnorm.py

@@ -1,4 +1,3 @@
-
 from keras import ops
 from keras.saving import register_keras_serializable
 
@@ -14,21 +13,22 @@ class ActNorm(InvertibleLayer):
     Activation Normalization is learned invertible normalization, using
     a Scale (s) and Bias (b) vector::
 
-       y = s * x + b (forward)
-       x = (y - b) / s (inverse)
+       y = s * x + b(forward)
+       x = (y - b) / s(inverse)
 
     References
     ----------
 
-    .. [1] Kingma, D. P., & Dhariwal, P. (2018). 
-        Glow: Generative flow with invertible 1x1 convolutions. 
+    .. [1] Kingma, D. P., & Dhariwal, P. (2018).
+        Glow: Generative flow with invertible 1x1 convolutions.
         Advances in Neural Information Processing Systems, 31.
 
     .. [2] Salimans, Tim, and Durk P. Kingma. (2016).
        Weight normalization: A simple reparameterization to accelerate
        training of deep neural networks.
        Advances in Neural Information Processing Systems, 29.
     """
+
     def __init__(self, **kwargs):
         super().__init__(**keras_kwargs(kwargs))
         self.scale = None