import threading
import warnings
from collections.abc import Callable, MutableSequence, Sequence
from typing import Any
import keras
from bayesflow._backend import jit
from bayesflow.types import Tensor
from bayesflow.utils.serialization import deserialize, serialize, serializable
from .transforms import (
AsCorrelationMatrix,
AsCovarianceMatrix,
AsSet,
AsSimplex,
AsTimeSeries,
Broadcast,
Concatenate,
Constrain,
ConvertDType,
Drop,
ExpandDims,
FilterTransform,
Group,
Keep,
Log,
MapTransform,
KerasTransform,
OneHot,
Rename,
SerializableCustomTransform,
Squeeze,
Sqrt,
Standardize,
ToArray,
Transform,
Ungroup,
RandomSubsample,
Take,
NanToNum,
)
from .transforms.filter_transform import Predicate
[docs]
@serializable("bayesflow.adapters")
class Adapter(MutableSequence[Transform]):
"""
Defines an adapter to apply various transforms to simulated / real
quantities in an amortized workflow. The role of the adapter is to
make simulation outputs / real data neural network-friendly.
Where possible, the transforms also supply an inverse transform.
An adapter can also have no transformers, in which case it simply
returns a shallow copy of the input data. This can be useful
for example to ensure that no side effects occur when the same data
dictionary is passed to downstream consumers.
Note: most ready-made transforms are rather simple and currently
support only numpy operations. However, the adapter is designed to
be flexible and can accommodate arbitrary Tensors (e.g., ``torch.Tensor``).
In the latter cases, users can implement custom transforms by inheriting
from the base :class:`~bayesflow.adapters.Transform` class and overriding
the ``forward`` and ``inverse`` methods.
Parameters
----------
transforms : Sequence[Transform], optional
The sequence of transforms to execute.
device: str, optional
Passed to keras.device to specify on which device to apply the adapter.
differentiable: bool, optional
If False (default), adapter transforms run in a state detached from the computational graph.
jit_compile: bool, optional
If True, lazily compile the transform sequence. The first call for each
direction and set of keyword arguments runs eagerly to initialize
stateful transforms. Dtype normalization and the final device handoff
remain outside the compiled region. Default is False.
"""
def __init__(
self,
transforms: Sequence[Transform] | None = None,
device: str = "cpu",
differentiable: bool = False,
jit_compile: bool = False,
):
if transforms is None:
transforms = []
self.transforms = list(transforms)
self.device = device
self.differentiable = differentiable
self.jit_compile = jit_compile
self._reset_compile_cache()
[docs]
@staticmethod
def create_default(inference_variables: Sequence[str]) -> "Adapter":
"""Create an adapter with a set of default transforms.
Parameters
----------
inference_variables : Sequence of str
The names of the variables to be inferred by an estimator.
Returns
-------
An initialized Adapter with a set of default transforms.
"""
return Adapter().concatenate(inference_variables, into="inference_variables")
[docs]
@classmethod
def from_config(cls, config: dict, custom_objects=None) -> "Adapter":
return cls(**deserialize(config, custom_objects=custom_objects))
[docs]
def get_config(self) -> dict:
config = {
"transforms": self.transforms,
"device": self.device,
"differentiable": self.differentiable,
"jit_compile": self.jit_compile,
}
return serialize(config)
[docs]
def forward(
self, data: dict[str, Any], *, log_det_jac: bool = False, **kwargs
) -> dict[str, Tensor] | tuple[dict[str, Tensor], dict[str, Tensor]]:
"""Apply the transforms in the forward direction.
Parameters
----------
data : dict[str, Any]
The data to be transformed.
log_det_jac: bool, optional
Whether to return the log determinant of the Jacobian of the transforms.
**kwargs : dict
Additional keyword arguments passed to each transform.
Returns
-------
dict | tuple[dict, dict]
The transformed data or tuple of transformed data and log determinant of the Jacobian.
"""
if not log_det_jac:
for transform in self.transforms:
data = transform(data, **kwargs)
return data
log_det_jac = {}
for transform in self.transforms:
transformed_data = transform(data, **kwargs)
log_det_jac = transform.log_det_jac(data, log_det_jac=log_det_jac, **kwargs)
data = transformed_data
return data, log_det_jac
[docs]
def inverse(
self, data: dict[str, Any], *, log_det_jac: bool = False, **kwargs
) -> dict[str, Tensor] | tuple[dict[str, Tensor], dict[str, Tensor]]:
"""Apply the transforms in the inverse direction.
Parameters
----------
data : dict[str, Any]
The data to be transformed.
log_det_jac: bool, optional
Whether to return the log determinant of the Jacobian of the transforms.
**kwargs : dict
Additional keyword arguments passed to each transform.
Returns
-------
dict | tuple[dict, dict]
The transformed data or tuple of transformed data and log determinant of the Jacobian.
"""
if not log_det_jac:
for transform in reversed(self.transforms):
data = transform(data, inverse=True, **kwargs)
return data
log_det_jac = {}
for transform in reversed(self.transforms):
data = transform(data, inverse=True, **kwargs)
log_det_jac = transform.log_det_jac(data, log_det_jac, inverse=True, **kwargs)
return data, log_det_jac
[docs]
def __call__(
self, data: dict[str, Any] | tuple[dict[str, Tensor], dict[str, Tensor]], *, inverse: bool = False, **kwargs
) -> dict[str, Tensor] | tuple[dict[str, Tensor], dict[str, Tensor]]:
"""Apply the transforms in the given direction.
Parameters
----------
data : Mapping[str, Any]
The data to be transformed.
inverse : bool, optional
If False, apply the forward transform, else apply the inverse transform (default False).
**kwargs
Additional keyword arguments passed to each transform.
Returns
-------
dict | tuple[dict, dict]
The transformed data or tuple of transformed data and log determinant of the Jacobian.
"""
# An identity adapter has no work to perform on its configured device.
# Go straight to the final handoff, which still normalizes dtype and
# preserves the differentiability contract on the default device.
if not self.transforms:
if kwargs.get("log_det_jac", False):
data = (data, {})
if isinstance(data, tuple):
return tuple(self.device_handoff(value) for value in data)
return self.device_handoff(data)
with keras.device(self.device):
data = keras.tree.map_structure(
lambda x: keras.ops.convert_to_tensor(x) if keras.ops.is_tensor(x) else x,
data,
)
data = self._apply_transforms(data, inverse=inverse, **kwargs)
if isinstance(data, tuple): # when log_det_jac=True,: (data, log_det_jac)
return tuple(self.device_handoff(value) for value in data)
return self.device_handoff(data)
def _apply_transforms(self, data, *, inverse: bool, **kwargs):
transform = self.inverse if inverse else self.forward
if not self.jit_compile:
return transform(data, **kwargs)
signature = (self.device, tuple(map(id, self.transforms)))
if signature != self._compiled_transform_signature:
self._reset_compile_cache(signature)
key = self._compile_key(inverse, kwargs)
if key is None or key in self._compile_failures:
return transform(data, **kwargs)
compiled = self._compiled_transforms.get(key)
if compiled is None:
# Several transforms initialize reversible metadata on their first
# eager call (e.g. Concatenate indices and ToArray input types).
# Warm that state before a compiler captures the transform graph.
result = transform(data, **kwargs)
try:
self._compiled_transforms[key] = jit(lambda value: transform(value, **kwargs))
except Exception as compile_error:
self._compile_failures.add(key)
self._warn_compile_failure(compile_error)
return result
if key not in self._compile_executed:
# Dataset workers can reach a newly-created lazy compiler at the
# same time. Serialize only this first compiler invocation; once
# built, steady-state calls remain parallel in the worker pool.
with self._compile_lock:
if key in self._compile_failures:
return transform(data, **kwargs)
if key not in self._compile_executed:
result = self._execute_compiled(compiled, transform, data, key, kwargs)
if key not in self._compile_failures:
self._compile_executed.add(key)
return result
return self._execute_compiled(compiled, transform, data, key, kwargs)
def _execute_compiled(self, compiled, transform, data, key, kwargs):
try:
return compiled(data)
except Exception as compile_error:
# Do not hide an actual adapter error behind compiler fallback. If
# eager execution also fails, surface that original transform error.
try:
result = transform(data, **kwargs)
except Exception:
raise
self._compile_failures.add(key)
self._warn_compile_failure(compile_error)
return result
@staticmethod
def _warn_compile_failure(compile_error: Exception) -> None:
warnings.warn(
f"Adapter JIT compilation failed on the {keras.backend.backend()!r} backend; "
f"falling back to eager transforms for this call signature. Compiler error: {compile_error}",
RuntimeWarning,
stacklevel=3,
)
@staticmethod
def _compile_key(inverse: bool, kwargs: dict) -> tuple | None:
"""Return a stable cache key for ordinary static transform kwargs."""
def freeze(value):
if value is None:
return ("none",)
if isinstance(value, bool):
return ("bool", value)
if isinstance(value, int):
return ("int", value)
if isinstance(value, float):
return ("float", value.hex())
if isinstance(value, str):
return ("str", value)
if isinstance(value, bytes):
return ("bytes", value)
if isinstance(value, tuple):
return ("tuple", tuple(freeze(item) for item in value))
if isinstance(value, list):
return ("list", tuple(freeze(item) for item in value))
if isinstance(value, dict):
return ("dict", tuple(sorted((name, freeze(item)) for name, item in value.items())))
raise TypeError
try:
return inverse, tuple(sorted((name, freeze(value)) for name, value in kwargs.items()))
except (TypeError, ValueError):
# Dynamic tensor/object kwargs should remain dynamic rather than be
# captured as constants in a compiled closure.
return None
def _reset_compile_cache(self, signature: tuple | None = None) -> None:
self._compiled_transforms = {}
self._compile_failures = set()
self._compile_executed = set()
self._compile_lock = threading.Lock()
self._compiled_transform_signature = (
(self.device, tuple(map(id, self.transforms))) if signature is None else signature
)
def __getstate__(self):
# Backend-compiled callables are process-local and generally not
# picklable. Dataset workers rebuild their own cache lazily.
state = self.__dict__.copy()
state.pop("_compiled_transforms", None)
state.pop("_compile_failures", None)
state.pop("_compile_executed", None)
state.pop("_compile_lock", None)
state.pop("_compiled_transform_signature", None)
return state
def __setstate__(self, state):
self.__dict__.update(state)
self._reset_compile_cache()
def __repr__(self):
result = ""
for i, transform in enumerate(self):
result += f"{i}: {transform!r}"
if i != len(self) - 1:
result += " -> "
return f"Adapter([{result}])"
[docs]
def append(self, value: Transform) -> "Adapter":
"""Append a transform to the list of transforms.
Parameters
----------
value : Transform
The transform to be added.
"""
self.transforms.append(value)
return self
[docs]
def device_handoff(self, data: tuple | dict[str, Any]) -> tuple | dict[str, Any]:
"""Convert ``data`` to ``floatx`` tensors and hand them off on the default device."""
floatx = keras.config.floatx()
def convert(x):
if keras.ops.is_tensor(x):
# cast in place: converting *with* a dtype moves to the target device first,
# which fails for dtypes that device does not support (float64 on torch MPS)
if keras.ops.dtype(x) != floatx:
x = keras.ops.cast(x, floatx)
else:
with keras.device(self.device):
x = keras.ops.convert_to_tensor(x, floatx)
# we need to hand off tensors on the default device, not on the device the adapter ran
x = keras.ops.convert_to_tensor(x)
return x if self.differentiable else keras.ops.stop_gradient(x)
return keras.tree.map_structure(convert, data)
def __delitem__(self, key: int | slice):
del self.transforms[key]
[docs]
def extend(self, values: Sequence[Transform]) -> "Adapter":
"""Extend the adapter with a sequence of transforms.
Parameters
----------
values : Sequence of Transform
The additional transforms to extend the adapter.
"""
if isinstance(values, Adapter):
values = values.transforms
self.transforms.extend(values)
return self
def __getitem__(self, item: int | slice) -> "Adapter":
if isinstance(item, int):
return self.transforms[item]
return Adapter(
self.transforms[item],
device=self.device,
differentiable=self.differentiable,
jit_compile=self.jit_compile,
)
[docs]
def insert(self, index: int, value: Transform | Sequence[Transform]) -> "Adapter":
"""Insert a transform at a given index.
Parameters
----------
index : int
The index to insert at.
value : Transform or Sequence of Transform
The transform or transforms to insert.
"""
if isinstance(value, Adapter):
value = value.transforms
if isinstance(value, Sequence):
# convenience: Adapters are always flat
self.transforms = self.transforms[:index] + list(value) + self.transforms[index:]
else:
self.transforms.insert(index, value)
return self
def __setitem__(self, key: int | slice, value: Transform | Sequence[Transform]) -> "Adapter":
if isinstance(value, Adapter):
value = value.transforms
if isinstance(key, int) and isinstance(value, Sequence):
if key < 0:
key += len(self.transforms)
key = slice(key, key + 1)
self.transforms[key] = value
return self
def __len__(self):
return len(self.transforms)
[docs]
def apply(
self,
include: str | Sequence[str] = None,
*,
forward: str,
inverse: str = None,
predicate: Predicate = None,
exclude: str | Sequence[str] = None,
**kwargs,
):
"""Append a :py:class:`~transforms.KerasTransform` to the adapter.
Parameters
----------
forward : str
The name of the NumPy function to use for the forward transformation.
inverse : str, optional
The name of the NumPy function to use for the inverse transformation.
By default, the inverse is inferred from the forward argument for supported methods.
You can find the supported methods in
:py:const:`~bayesflow.adapters.transforms.KerasTransform.INVERSE_METHODS`.
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=KerasTransform,
predicate=predicate,
include=include,
exclude=exclude,
forward=forward,
inverse=inverse,
**kwargs,
)
self.transforms.append(transform)
return self
[docs]
def apply_serializable(
self,
include: str | Sequence[str] = None,
*,
forward: Callable[[Tensor], Tensor],
inverse: Callable[[Tensor], Tensor],
predicate: Predicate = None,
exclude: str | Sequence[str] = None,
**kwargs,
):
"""Append a :py:class:`~transforms.SerializableCustomTransform` to the adapter.
Parameters
----------
forward : function, no lambda
Registered serializable 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 can be used here.
inverse : function, no lambda
Registered serializable 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
functions can be used here.
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.
Raises
------
ValueError
When the provided functions are not registered serializable functions.
Notes
-----
Important: The forward and inverse functions have to be registered with Keras.
To do so, use the `@keras.saving.register_keras_serializable` decorator.
They must also be registered (and identical) when loading the adapter
at a later point in time.
Examples
--------
The example below shows how to use the
`keras.saving.register_keras_serializable` decorator to
register functions with Keras. Note that for this simple
example, one usually would use the simpler :py:meth:`apply`
method.
>>> import keras
>>>
>>> @keras.saving.register_keras_serializable("custom")
>>> def forward_fn(x):
>>> return x**2
>>>
>>> @keras.saving.register_keras_serializable("custom")
>>> def inverse_fn(x):
>>> return x**0.5
>>>
>>> adapter = bf.Adapter().apply_serializable(
>>> "x",
>>> forward=forward_fn,
>>> inverse=inverse_fn,
>>> )
"""
transform = FilterTransform(
transform_constructor=SerializableCustomTransform,
predicate=predicate,
include=include,
exclude=exclude,
forward=forward,
inverse=inverse,
**kwargs,
)
self.transforms.append(transform)
return self
[docs]
def as_set(self, keys: str | Sequence[str]):
"""Append an :py:class:`~transforms.AsSet` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to apply the transform to.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: AsSet() for key in keys})
self.transforms.append(transform)
return self
[docs]
def as_time_series(self, keys: str | Sequence[str]):
"""Append an :py:class:`~transforms.AsTimeSeries` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to apply the transform to.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: AsTimeSeries() for key in keys})
self.transforms.append(transform)
return self
[docs]
def broadcast(
self,
keys: str | Sequence[str],
*,
to: str,
expand: str | int | tuple = "left",
exclude: int | tuple = -1,
squeeze: int | tuple = None,
):
"""Append a :py:class:`~transforms.Broadcast` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to apply the transform to.
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
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.
Notes
-----
Important: Do not broadcast to variables that are used as inference variables
(i.e., parameters to be inferred by the networks). The adapter will work during training
but then fail during inference because the variable being broadcasted to is not available.
"""
if isinstance(keys, str):
keys = [keys]
transform = Broadcast(keys, to=to, expand=expand, exclude=exclude, squeeze=squeeze)
self.transforms.append(transform)
return self
[docs]
def clear(self):
"""Remove all transforms from the adapter."""
self.transforms = []
return self
[docs]
def concatenate(self, keys: str | Sequence[str], *, into: str, axis: int = -1):
"""Append a :py:class:`~transforms.Concatenate` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to concatenate.
into : str
The name of the resulting variable.
axis : int, optional
Along which axis to concatenate the keys. The last axis is used by default.
"""
if isinstance(keys, str):
transform = Rename(keys, to_key=into)
else:
transform = Concatenate(keys, into=into, axis=axis)
self.transforms.append(transform)
return self
[docs]
def convert_dtype(
self,
from_dtype: str,
to_dtype: str,
*,
predicate: Predicate = None,
include: str | Sequence[str] = None,
exclude: str | Sequence[str] = None,
):
"""Append a :py:class:`~transforms.ConvertDType` transform to the adapter.
See also :py:meth:`~bayesflow.adapters.Adapter.map_dtype`.
Parameters
----------
from_dtype : str
Original dtype
to_dtype : str
Target dtype
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.
"""
transform = FilterTransform(
transform_constructor=ConvertDType,
predicate=predicate,
include=include,
exclude=exclude,
from_dtype=from_dtype,
to_dtype=to_dtype,
)
self.transforms.append(transform)
return self
[docs]
def constrain(
self,
keys: str | Sequence[str],
*,
lower: int | float | Tensor = None,
upper: int | float | Tensor = None,
method: str = "default",
inclusive: str = "both",
epsilon: float = 1e-6,
):
"""Append a :py:class:`~transforms.Constrain` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to constrain.
lower: int or float or Tensor, optional
Lower bound for named data variable.
upper : int or float or Tensor, optional
Upper bound for named data variable.
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
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
Small value to ensure inclusive bounds are not violated.
Current default is 1e-6 as this ensures finite outcomes
with the default transformations applied to data exactly at the boundaries.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform(
transform_map={
key: Constrain(lower=lower, upper=upper, method=method, inclusive=inclusive, epsilon=epsilon)
for key in keys
}
)
self.transforms.append(transform)
return self
[docs]
def as_correlation_matrix(self, keys: str | Sequence[str], cholesky: bool = False):
"""Append a :py:class:`~transforms.AsCorrelationMatrix` transform to the adapter.
Constrains neural network predictions of a data variable to a valid (symmetric, positive definite,
unit diagonal) K x K correlation matrix, using the transforms explained by [1].
The unconstrained representation is a flat vector of K * (K - 1) / 2 entries.
Parameters
----------
keys : str or Sequence of str
The names of the variables to constrain.
cholesky : bool, optional
Whether the *constrained* side of the transform is the lower Cholesky factor,
rather than the full correlation matrix.
Default is False (use the full correlation matrix).
References
----------
[1] Lewandowski, D., Kurowicka, D., & Joe, H. (2009).
Generating random correlation matrices based on vines and extended onion method.
Journal of Multivariate Analysis, 100(9), 1989-2001. https://doi.org/10.1016/j.jmva.2009.04.008
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform(transform_map={key: AsCorrelationMatrix(cholesky=cholesky) for key in keys})
self.transforms.append(transform)
return self
[docs]
def as_covariance_matrix(self, keys: str | Sequence[str], cholesky: bool = False, diag_kwargs: dict = None):
"""Append a :py:class:`~transforms.AsCovarianceMatrix` transform to the adapter.
Constrains neural network predictions of a data variable to a valid (symmetric,
positive definite) K x K covariance (or precision) matrix.
The unconstrained representation is a flat vector of K * (K + 1) / 2 entries
of the lower Cholesky factor of the covariance matrix.
First, the K diagonal entries, passed through a lower-bounded :py:class:`~transforms.Constrain`
transform to keep them positive, followed by the K * (K - 1) / 2 entries from the lower triangular.
Parameters
----------
keys : str or Sequence of str
The names of the variables to constrain.
cholesky : bool, optional
Whether the *constrained* side of the transform is the Cholesky factor,
rather than the full covariance matrix.
Default is False (use the full covariance matrix).
diag_kwargs : dict, optional
Keyword arguments forwarded to the :py:class:`~transforms.Constrain` transform
that constrains the diagonal to be positive.
The `lower` bound is always fixed to 0.0.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform(
transform_map={key: AsCovarianceMatrix(cholesky=cholesky, diag_kwargs=diag_kwargs) for key in keys}
)
self.transforms.append(transform)
return self
[docs]
def as_precision_matrix(self, **kwargs):
"""Append a :py:class:`~transforms.AsCovarianceMatrix` transform to the adapter.
A precision matrix is the inverse of a covariance matrix, but is likewise symmetric
and positive definite, so the same transform applies.
This is an alias for :py:meth:`~Adapter.as_covariance_matrix`;
see there for the accepted parameters.
Note that this transform does not itself invert a covariance matrix into a
precision matrix (or vice versa), it only constrains its input to *some* symmetric
positive definite matrix, whatever you intend that matrix to represent.
"""
return self.as_covariance_matrix(**kwargs)
[docs]
def as_simplex(self, keys: str | Sequence[str], axis: int | Tensor = -1, method: str = "default"):
"""Append a :py:class:`~transforms.AsSimplex` transform to the adapter.
Constrains neural network predictions of a data variable to a unit simplex, so that the
constrained representation is non-negative and sums to one along the given axis.
Parameters
----------
keys : str or Sequence of str
The names of the variables to constrain.
axis : int, optional
The axis of the *simplex-constrained* data along which values sum to one over all K elements.
The corresponding unconstrained representation has a size of K-1 along this axis.
method : str, optional
Method by which to transform between the K dimensional simplex space
and the K-1 dimensional unconstrained space.
- "default" / "simplex": sum-to-zero (orthogonal basis) transform followed by a softmax transform.
Plain softmax is not invertible (shifting all values by a constant results in the same inverse).
The sum-to-zero removes this issue.
- "stick": stick-breaking logistic transform.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform(transform_map={key: AsSimplex(axis=axis, method=method) for key in keys})
self.transforms.append(transform)
return self
[docs]
def drop(self, keys: str | Sequence[str]):
"""Append a :py:class:`~transforms.Drop` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to drop.
"""
if isinstance(keys, str):
keys = [keys]
transform = Drop(keys)
self.transforms.append(transform)
return self
[docs]
def expand_dims(self, keys: str | Sequence[str], *, axis: int | tuple):
"""Append an :py:class:`~transforms.ExpandDims` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to expand.
axis : int or tuple
The axis to expand.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: ExpandDims(axis=axis) for key in keys})
self.transforms.append(transform)
return self
[docs]
def group(self, keys: Sequence[str], into: str, *, prefix: str = ""):
"""Append a :py:class:`~transforms.Group` transform to the adapter.
Groups the given variables as a dictionary in the key `into`. As most transforms do
not support nested structures, this should usually be the last transform in the adapter.
Parameters
----------
keys : Sequence of str
The names of the variables to group together.
into : str
The name of the variable to store the grouped variables in.
prefix : str, optional
An optional common prefix of the variable names before grouping, which will be removed after grouping.
Raises
------
ValueError
If a prefix is specified, but a provided key does not start with the prefix.
"""
if isinstance(keys, str):
keys = [keys]
transform = Group(keys=keys, into=into, prefix=prefix)
self.transforms.append(transform)
return self
[docs]
def ungroup(self, key: str, *, prefix: str = ""):
"""Append an :py:class:`~transforms.Ungroup` transform to the adapter.
Ungroups the the variables in `key` from a dictionary into individual entries. Most transforms do
not support nested structures, so this can be used to flatten a nested structure.
The nesting can be re-established after the transforms using the :py:meth:`group` method.
Parameters
----------
key : str
The name of the variable to ungroup. The corresponding variable has to be a dictionary.
prefix : str, optional
An optional common prefix that will be added to the ungrouped variable names. This can be necessary
to avoid duplicate names.
"""
transform = Ungroup(key=key, prefix=prefix)
self.transforms.append(transform)
return self
[docs]
def keep(self, keys: str | Sequence[str]):
"""Append a :py:class:`~transforms.Keep` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to keep.
"""
if isinstance(keys, str):
keys = [keys]
transform = Keep(keys)
self.transforms.append(transform)
return self
[docs]
def log(self, keys: str | Sequence[str], *, p1: bool = False):
"""Append an :py:class:`~transforms.Log` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to transform.
p1 : boolean
Add 1 to the input before taking the logarithm?
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: Log(p1=p1) for key in keys})
self.transforms.append(transform)
return self
[docs]
def map_dtype(self, keys: str | Sequence[str], to_dtype: str):
"""Append a :py:class:`~transforms.ConvertDType` transform to the adapter.
See also :py:meth:`~bayesflow.adapters.Adapter.convert_dtype`.
Parameters
----------
keys : str or Sequence of str
The names of the variables to transform.
to_dtype : str
Target dtype
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: ConvertDType(to_dtype) for key in keys})
self.transforms.append(transform)
return self
[docs]
def one_hot(self, keys: str | Sequence[str], num_classes: int):
"""Append a :py:class:`~transforms.OneHot` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to transform.
num_classes : int
Number of classes for the encoding.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: OneHot(num_classes=num_classes) for key in keys})
self.transforms.append(transform)
return self
[docs]
def random_subsample(self, key: str, *, sample_size: int | float, axis: int = -1):
"""
Append a :py:class:`~transforms.RandomSubsample` transform to the adapter.
Parameters
----------
key : str or Sequence of str
The name of the variable to subsample.
sample_size : int or float
The number of samples to draw, or a fraction between 0 and 1 of the total number of samples to draw.
axis: int, optional
Which axis to draw samples over. The last axis is used by default.
"""
if not isinstance(key, str):
raise TypeError("Can only subsample one batch entry at a time.")
transform = MapTransform({key: RandomSubsample(sample_size=sample_size, axis=axis)})
self.transforms.append(transform)
return self
[docs]
def rename(self, from_key: str, to_key: str):
"""Append a :py:class:`~transforms.Rename` transform to the adapter.
Parameters
----------
from_key : str
Variable name that should be renamed
to_key : str
New variable name
"""
self.transforms.append(Rename(from_key, to_key))
return self
[docs]
def scale(self, keys: str | Sequence[str], by: float | Tensor):
from .transforms import Scale
if isinstance(keys, str):
keys = [keys]
self.transforms.append(MapTransform({key: Scale(scale=by) for key in keys}))
return self
[docs]
def shift(self, keys: str | Sequence[str], by: float | Tensor):
from .transforms import Shift
if isinstance(keys, str):
keys = [keys]
self.transforms.append(MapTransform({key: Shift(shift=by) for key in keys}))
return self
[docs]
def split(self, key: str, *, into: Sequence[str], indices_or_sections: int | Sequence[int] = None, axis: int = -1):
from .transforms import Split
if isinstance(into, str):
transform = Rename(key, into)
else:
transform = Split(key, into, indices_or_sections, axis)
self.transforms.append(transform)
return self
[docs]
def squeeze(self, keys: str | Sequence[str], *, axis: int | Sequence[int]):
"""Append a :py:class:`~transforms.Squeeze` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to squeeze.
axis : int or tuple
The axis to squeeze. As the number of batch dimensions might change, we advise using negative
numbers (i.e., indexing from the end instead of the start).
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: Squeeze(axis=axis) for key in keys})
self.transforms.append(transform)
return self
[docs]
def sqrt(self, keys: str | Sequence[str]):
"""Append an :py:class:`~transforms.Sqrt` transform to the adapter.
Parameters
----------
keys : str or Sequence of str
The names of the variables to transform.
"""
if isinstance(keys, str):
keys = [keys]
transform = MapTransform({key: Sqrt() for key in keys})
self.transforms.append(transform)
return self
[docs]
def standardize(
self,
include: str | Sequence[str] = None,
*,
mean: int | float | Tensor,
std: int | float | Tensor,
predicate: Predicate = None,
exclude: str | Sequence[str] = None,
**kwargs,
):
"""Append a :py:class:`~transforms.Standardize` transform to the adapter.
Parameters
----------
include : str or Sequence of str, optional
Names of variables to include in the transform.
mean : int or float
Specifies the mean (location) of the transform.
std : int or float
Specifies the standard deviation (scale) of the transform.
predicate : Predicate, optional
Function that indicates which variables should be transformed.
exclude : str or Sequence of str, optional
Names of variables to exclude from the transform.
**kwargs :
Additional keyword arguments passed to the transform.
"""
transform = FilterTransform(
transform_constructor=Standardize,
predicate=predicate,
include=include,
exclude=exclude,
mean=mean,
std=std,
**kwargs,
)
self.transforms.append(transform)
return self
[docs]
def take(
self,
include: str | Sequence[str] = None,
*,
indices: Sequence[int],
axis: int = -1,
predicate: Predicate = None,
exclude: str | Sequence[str] = None,
):
"""
Append a :py:class:`~transforms.Take` transform to the adapter.
Parameters
----------
include : str or Sequence of str, optional
Names of variables to include in the transform.
indices : Sequence of int
Which indices to take from the data.
axis : int, optional
Which axis to take from. The last axis is used by default.
predicate : Predicate, optional
Function that indicates which variables should be transformed.
exclude : str or Sequence of str, optional
Names of variables to exclude from the transform.
"""
transform = FilterTransform(
transform_constructor=Take,
predicate=predicate,
include=include,
exclude=exclude,
indices=indices,
axis=axis,
)
self.transforms.append(transform)
return self
[docs]
def to_array(
self,
include: str | Sequence[str] = None,
*,
predicate: Predicate = None,
exclude: str | Sequence[str] = None,
**kwargs,
):
"""Append a :py:class:`~transforms.ToArray` 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=ToArray,
predicate=predicate,
include=include,
exclude=exclude,
**kwargs,
)
self.transforms.append(transform)
return self
[docs]
def to_dict(self):
from .transforms import ToDict
transform = ToDict()
self.transforms.append(transform)
return self
[docs]
def nan_to_num(
self,
keys: str | Sequence[str],
default_value: float = 0.0,
return_mask: bool = False,
mask_prefix: str = "mask",
):
"""
Append :py:class:`~bf.adapters.transforms.NanToNum` transform to the adapter.
Parameters
----------
keys : str or sequence of str
The names of the variables to clean / mask.
default_value : float
Value to substitute wherever data is NaN. Defaults to 0.0.
return_mask : bool
If True, encode a binary missingness mask alongside the data, which is also used to restore the
NaNs in the inverse transform. If False, the transform is not invertible and the inverse leaves
the data unchanged. Defaults to False.
mask_prefix : str
Prefix for the mask key in the output dictionary. Defaults to 'mask_'. If the mask key already exists,
a ValueError is raised to avoid overwriting existing masks.
"""
if isinstance(keys, str):
keys = [keys]
for key in keys:
self.transforms.append(
NanToNum(key=key, default_value=default_value, return_mask=return_mask, mask_prefix=mask_prefix)
)
return self
