Subset arrays

.gitignore

@@ -39,3 +39,6 @@ docs/
 
 # MacOS
 .DS_Store
+
+# Rproj
+.Rproj.user

bayesflow/adapters/adapter.py

@@ -27,6 +27,8 @@
     Standardize,
     ToArray,
     Transform,
+    SubsampleArray,
+    Take
 )
 from .transforms.filter_transform import Predicate
 
@@ -541,6 +543,42 @@ def one_hot(self, keys: str | Sequence[str], num_classes: int):
         transform = MapTransform({key: OneHot(num_classes=num_classes) for key in keys})
         self.transforms.append(transform)
         return self
+
+    def random_subsample(self,
+        keys: str | Sequence[str],
+        *,
+        sample_size: int,
+        axis: int=-1,
+        **kwargs,
+    ):
+        """
+        Append a :py:class:`~transforms.SubsampleArray` transform to the adapter.
+
+        Parameters
+        ----------
+        predicate : Predicate, optional
+            Function that indicates which variables should be transformed.
+        include : str or Sequence of str, optional
+            Names of variables to include in the transform.
+        exclude : str or Sequence of str, optional
+            Names of variables to exclude from the transform.
+        **kwargs : dict
+            Additional keyword arguments passed to the transform.
+
+        """
+        if isinstance(keys, str):
+            keys = [keys]
+
+        transform = MapTransform(
+            transform_map={
+                key: SubsampleArray(sample_size=sample_size, axis=axis)
+                for key in keys
+            }
+
+        )
+
+        self.transforms.append(transform)
+        return self
 
     def rename(self, from_key: str, to_key: str):
         """Append a :py:class:`~transforms.Rename` transform to the adapter.
@@ -601,6 +639,36 @@ def standardize(
         self.transforms.append(transform)
         return self
 
+    def take(self,
+        *,
+        predicate: Predicate = None,
+        include: str | Sequence[str] = None,
+        exclude: str | Sequence[str] = None,
+        **kwargs,):
+        """
+        Append a :py:class:`~transforms.Take` transform to the adapter.
+
+        Parameters
+        ----------
+        predicate : Predicate, optional
+            Function that indicates which variables should be transformed.
+        include : str or Sequence of str, optional
+            Names of variables to include in the transform.
+        exclude : str or Sequence of str, optional
+            Names of variables to exclude from the transform.
+        **kwargs : dict
+            Additional keyword arguments passed to the transform. """
+        transform = FilterTransform(
+            transform_constructor=Take,
+            predicate=predicate,
+            include=include,
+            exclude=exclude,
+            **kwargs,
+        )
+        self.transforms.append(transform)
+        return self
+
+
     def to_array(
         self,
         *,
@@ -631,3 +699,9 @@ def to_array(
         )
         self.transforms.append(transform)
         return self
+
+
+
+
+
+

bayesflow/adapters/transforms/__init__.py

@@ -18,6 +18,8 @@
 from .standardize import Standardize
 from .to_array import ToArray
 from .transform import Transform
+from .subsample_array import SubsampleArray
+from .take import Take
 
 from ...utils._docs import _add_imports_to_all
 

bayesflow/adapters/transforms/subsample_array.py

@@ -0,0 +1,39 @@
+import numpy as np
+from keras.saving import register_keras_serializable as serializable
+
+from .elementwise_transform import ElementwiseTransform
+
+
+@serializable(package="bayesflow.adapters")
+class SubsampleArray(ElementwiseTransform):
+    """
+    A transform that takes a random subsample of the data within an axis.
+
+    Example: adapter.subsample("x", sample_size = 3, axis = -1)
+
+    """
+
+    def __init__(
+            self,
+            sample_size: int, 
+            axis: int = -1, 
+                 ):
+        super().__init__()
+        self.sample_size = sample_size
+        self.axis = axis 
+
+    def forward(self, data: np.ndarray):
+        sample_size = self.sample_size
+        axis = self.axis 
+
+        max_sample_size = data.shape[axis]
+
+        sample_indices = np.random.permutation(max_sample_size)[
+            0 : sample_size - 1
+        ]  # random sample without replacement
+
+        return np.take(data, sample_indices, axis)
+
+    def inverse(self, data, **kwargs):
+        # non invertible transform
+        return data

bayesflow/adapters/transforms/take.py

@@ -0,0 +1,24 @@
+import numpy as np
+from keras.saving import register_keras_serializable as serializable
+
+from .elementwise_transform import ElementwiseTransform
+
+
+@serializable(package="bayesflow.adapters")
+class Take(ElementwiseTransform):
+    """
+    A transform to reduce the dimensionality of arrays output by the summary network
+    Axis is a mandatory argument and will default  to the last axis.
+    Example: adapter.take("x", np.arange(0,3), axis=-1)
+
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, data, indices, axis=-1):
+        return np.take(data, indices, axis)
+
+    def inverse(self, data):
+        # not a true invertible function
+        return data

tests/test_adapters/conftest.py

@@ -25,6 +25,8 @@ def adapter():
         .one_hot("o1", 10)
         .keep(["x", "y", "z1", "p1", "p2", "s1", "s2", "t1", "t2", "o1"])
         .rename("o1", "o2")
+        .random_subsample("s3", sample_size = 33, axis = 0)
+        .take("s3", indices = np.arange(0,32), axis = 0)
     )
 
     return d
@@ -47,4 +49,5 @@ def random_data():
         "d1": np.random.standard_normal(size=(32, 2)),
         "d2": np.random.standard_normal(size=(32, 2)),
         "o1": np.random.randint(0, 9, size=(32, 2)),
+        "s3": np.random.standard_normal(size=(35,2))
     }

tests/test_links/test_links.py

@@ -30,9 +30,9 @@ def check_ordering(output, axis):
     assert np.all(np.diff(output, axis=axis) > 0), f"is not ordered along specified axis: {axis}."
     for i in range(output.ndim):
         if i != axis % output.ndim:
-            assert not np.all(np.diff(output, axis=i) > 0), (
-                f"is ordered along axis which is not meant to be ordered: {i}."
-            )
+            assert not np.all(
+                np.diff(output, axis=i) > 0
+            ), f"is ordered along axis which is not meant to be ordered: {i}."
 
 
 @pytest.mark.parametrize("axis", [0, 1, 2])
@@ -69,6 +69,6 @@ def test_positive_semi_definite(random_matrix_batch):
     output = keras.ops.convert_to_numpy(output)
     eigenvalues = np.linalg.eig(output).eigenvalues
 
-    assert np.all(eigenvalues.real > 0) and np.all(np.isclose(eigenvalues.imag, 0)), (
-        f"output is not positive semi-definite: real={eigenvalues.real}, imag={eigenvalues.imag}"
-    )
+    assert np.all(eigenvalues.real > 0) and np.all(
+        np.isclose(eigenvalues.imag, 0)
+    ), f"output is not positive semi-definite: real={eigenvalues.real}, imag={eigenvalues.imag}"

tests/test_networks/test_point_inference_network/test_point_inference_network.py

@@ -44,9 +44,9 @@ def test_save_and_load(tmp_path, point_inference_network, random_samples, random
 
     for key_outer in out1.keys():
         for key_inner in out1[key_outer].keys():
-            assert keras.ops.all(keras.ops.isclose(out1[key_outer][key_inner], out2[key_outer][key_inner])), (
-                "Output of original and loaded model differs significantly."
-            )
+            assert keras.ops.all(
+                keras.ops.isclose(out1[key_outer][key_inner], out2[key_outer][key_inner])
+            ), "Output of original and loaded model differs significantly."
 
 
 def test_copy_unequal(point_inference_network, random_samples, random_conditions):