examples for keep and to_array in style of broadcast as well as usage

Author: eodole

bayesflow/adapters/transforms/keep.py

@@ -11,14 +11,15 @@
 
 @serializable(package="bayesflow.adapters")
 class Keep(Transform):
-    '''
+    """
     Name the data parameters that should be kept for futher calculation. 
 
     Parameters: 
 
         cls: tuple containing the names of kept data variables as strings. 
 
-    Example: 
+    Example 1: 
+
         Two moons simulator generates data for priors alpha, r and theta as well as observation data x.
         We are interested only in theta and x, to keep only theta and x we should use the following; 
 
@@ -30,8 +31,20 @@ class Keep(Transform):
 
             )
 
+    Example 2: 
+
+    >>> a = [1,2,3,4]
+    >>> b = [[1,2],[3,4]]
+    >>> c = [[5,6,7,8]]
+    >>> dat = dict(a=a,b=b,c =c)
+    # Here we want to only keep elements b and c 
+    >>> keeper = bf.adapters.transforms.Keep(("b","c"))
+    >>> keeper.forward(dat)
+    {'b': [[1, 2], [3, 4]], 'c': [[5, 6, 7, 8]]}
+
+
 
-    '''
+    """
     def __init__(self, keys: Sequence[str]):
         self.keys = keys
 

bayesflow/adapters/transforms/to_array.py

@@ -11,6 +11,22 @@
 
 @serializable(package="bayesflow.adapters")
 class ToArray(ElementwiseTransform):
+    """
+    Checks provided data for any non-arrays and converts them to numpy arrays.
+    This ensures all data is in a format suitable for training. 
+
+    Example: 
+    >>> ta = bf.adapters.transforms.ToArray()
+    >>> a = [1,2,3,4]
+    >>> ta.forward(a)
+        array([1, 2, 3, 4])
+    >>> b = [[1,2],[3,4]]
+    >>> ta.forward(b)
+        array([[1, 2],
+            [3, 4]])
+    """
+
+
     def __init__(self):
         super().__init__()
         self.original_type = None