[no ci] add docstrings to Adapter, minor formatting

Part of #340.

Author: vpratz

bayesflow/adapters/adapter.py

@@ -17,18 +17,21 @@ class Broadcast(Transform):
 
     Parameters
     ----------
-    keys: sequence of str,
+    keys : sequence of str,
         Input a list of strings, where the strings are the names of data variables.
-    expand: str or int or tuple, optional
+    to : str
+        Name of the data variable to broadcast to.
+    expand : str or int or tuple, optional
         Where should new dimensions be added to match the number of dimensions in `to`?
         Can be "left", "right", or an integer or tuple containing the indices of the new dimensions.
         The latter is needed if we want to include a dimension in the middle, which will be required
         for more advanced cases. By default we expand left.
-
-    exclude: int or tuple, optional
+    exclude : int or tuple, optional
         Which dimensions (of the dimensions after expansion) should retain their size,
         rather than being broadcasted to the corresponding dimension size of `to`?
         By default we exclude the last dimension (usually the data dimension) from broadcasting the size.
+    squeeze : int or tuple, optional
+        Axis to squeeze after broadcasting.
 
     Examples
     --------

bayesflow/adapters/transforms/broadcast.py

@@ -16,12 +16,12 @@ class Concatenate(Transform):
 
     Parameters
     ----------
-    keys: sequence of str,
+    keys : sequence of str,
         Input a list of strings, where the strings are the names of data variables.
-    into: str
+    into : str
         A string telling the network how to use the variables named in keys.
-    axis: int, optional
-        integer specifing along which axis to concatonate the keys. The last axis is used by default.
+    axis : int, optional
+        Along which axis to concatenate the keys. The last axis is used by default.
 
     Examples
     --------

bayesflow/adapters/transforms/concatenate.py

@@ -23,22 +23,22 @@ class Constrain(ElementwiseTransform):
     * : str
         String containing the name of the data variable to be transformed e.g. "sigma". See examples below.
 
-    lower: int or float or np.darray, optional
+    lower : int or float or np.darray, optional
         Lower bound for named data variable.
-    upper: int or float or np.darray, optional
+    upper : int or float or np.darray, optional
         Upper bound for named data variable.
-    method: str, optional
+    method : str, optional
         Method by which to shrink the network predictions space to specified bounds. Choose from
         - Double bounded methods: sigmoid, expit, (default = sigmoid)
         - Lower bound only methods: softplus, exp, (default = softplus)
         - Upper bound only methods: softplus, exp, (default = softplus)
-    inclusive: {'both', 'lower', 'upper', 'none'}, optional
+    inclusive : {'both', 'lower', 'upper', 'none'}, optional
         Indicates which bounds are inclusive (or exclusive).
         - "both" (default): Both lower and upper bounds are inclusive.
         - "lower": Lower bound is inclusive, upper bound is exclusive.
         - "upper": Lower bound is exclusive, upper bound is inclusive.
         - "none": Both lower and upper bounds are exclusive.
-    epsilon: float, optional
+    epsilon : float, optional
         Small value to ensure inclusive bounds are not violated.
         Current default is 1e-15 as this ensures finite outcomes
         with the default transformations applied to data exactly at the boundaries.

bayesflow/adapters/transforms/constrain.py

@@ -16,7 +16,7 @@ class Drop(Transform):
 
     Parameters
     ----------
-    keys: sequence of str
+    keys : sequence of str
         Names of data variables that should be dropped
 
     Examples

bayesflow/adapters/transforms/drop.py

@@ -13,6 +13,13 @@ class ExpandDims(ElementwiseTransform):
     """
     Expand the shape of an array.
 
+    Parameters
+    ----------
+    keys : str or Sequence of str
+        The names of the variables to expand.
+    axis : int or tuple
+        The axis to expand.
+
     Examples
     --------
     shape (3,) array:

bayesflow/adapters/transforms/expand_dims.py

@@ -16,7 +16,7 @@ class Keep(Transform):
 
     Parameters
     ----------
-    keys: sequence of str
+    keys : sequence of str
         The names of kept data variables as strings.
 
     Examples

bayesflow/adapters/transforms/keep.py

@@ -16,12 +16,12 @@ class LambdaTransform(ElementwiseTransform):
 
     Parameters
     ----------
-    forward: callable, no lambda
+    forward : callable, no lambda
         Function to transform the data in the forward pass.
         For the adapter to be serializable, this function has to be serializable
         as well (see Notes). Therefore, only proper functions and no lambda
         functions should be used here.
-    inverse: callable, no lambda
+    inverse : callable, no lambda
         Function to transform the data in the inverse pass.
         For the adapter to be serializable, this function has to be serializable
         as well (see Notes). Therefore, only proper functions and no lambda

bayesflow/adapters/transforms/lambda_transform.py

@@ -13,9 +13,9 @@ class Rename(Transform):
 
     Parameters
     ----------
-    from_key: str
+    from_key : str
         Variable name that should be renamed
-    to_key: str
+    to_key : str
         New variable name
 
     Examples

bayesflow/adapters/transforms/rename.py

@@ -20,12 +20,12 @@ class Standardize(ElementwiseTransform):
     ----------
     mean : int or float, optional
         Specify a mean if known but will be estimated from data when not provided
-    std: int or float, optional
+    std : int or float, optional
         Specify a standard devation if known but will be estimated from data when not provided
-    axis: int, optional
+    axis : int, optional
         A specific axis along which standardization should take place. By default
         standardization happens individually for each dimension
-    momentum: float in (0,1)
+    momentum : float in (0,1)
         The momentum during training
     """