from typing import Sequence, Mapping, Any, Literal
import numpy as np
import keras
from bayesflow.utils.exceptions import ShapeError
from bayesflow.networks import MLP
[docs]
def classifier_two_sample_test(
estimates: np.ndarray,
targets: np.ndarray,
metric: str = "accuracy",
patience: int = 5,
max_epochs: int = 1000,
batch_size: int = 128,
return_metric_only: bool = True,
cross_validation_splits: int = 5,
validation_split: float = 0.5,
standardize: bool = True,
mlp_widths: Sequence | Literal["auto"] = "auto",
**kwargs,
) -> float | Mapping[str, Any]:
"""
C2ST metric [1] between samples from two distributions computed using a neural classifier.
Can be computationally expensive if called in a loop[, since it needs to train the model
for each set of samples.
Note: works best for large numbers of samples and averaged across different posteriors.
[1] Lopez-Paz, D., & Oquab, M. (2016). Revisiting classifier two-sample tests. arXiv:1610.06545.
Parameters
----------
estimates : np.ndarray
Array of shape (num_samples_est, num_variables) containing samples representing estimated quantities
(e.g., approximate posterior samples).
targets : np.ndarray
Array of shape (num_samples_tar, num_variables) containing target samples
(e.g., samples from a reference posterior).
metric : str, optional
Metric to evaluate the classifier performance. Default is "accuracy".
patience : int, optional
Number of epochs with no improvement after which training will be stopped. Default is 5.
max_epochs : int, optional
Maximum number of epochs to train the classifier. Default is 1000.
batch_size : int, optional
Number of samples per batch during training. Default is 64.
return_metric_only : bool, optional
If True, only the final validation metric is returned. Otherwise, a dictionary with the score, classifier, and
full training history is returned. Default is True.
cross_validation_splits : int, optional
Number of cross-validation splits to perform. Default is 5.
validation_split : float, optional
Fraction of the training data to be used as validation data, for single hold-out split. Default is 0.5.
standardize : bool, optional
If True, both estimates and targets will be standardized using the mean and standard deviation of estimates.
Default is True.
mlp_widths : Sequence[int], optional
Sequence specifying the number of units in each hidden layer of the MLP classifier.
If set to 'auto', defaults to two hidden layers with widths such that width is larger than 10 times
the number of variables and a power of two. Default is 'auto'.
**kwargs
Additional keyword arguments. Recognized keyword:
mlp_kwargs : dict
Dictionary of additional parameters to pass to the MLP constructor.
Returns
-------
results : float or dict
If return_metric_only is True, returns the final validation metric (e.g., accuracy) as a float.
Otherwise, returns a dictionary with keys "score", "classifier", and "history", where "score"
is the final validation metric, "classifier" is the trained Keras model, and "history" contains the
full training history.
"""
# Error, if targets dim does not match estimates dim
num_dims = estimates.shape[1]
if not num_dims == targets.shape[1]:
raise ShapeError(
f"estimates and targets can have different number of samples (1st dim)"
f"but must have the same dimensionality (2nd dim)"
f"found: estimates shape {estimates.shape[1]}, targets shape {targets.shape[1]}"
)
if mlp_widths == "auto":
widths = 2 ** int(np.ceil(np.log2(10 * num_dims)))
mlp_widths = [widths, widths]
# Standardize both estimates and targets relative to estimates mean and std
if standardize:
estimates_mean = np.mean(estimates, axis=0)
estimates_std = np.std(estimates, axis=0)
estimates = (estimates - estimates_mean) / estimates_std
targets = (targets - estimates_mean) / estimates_std
# Create data for classification task
data = np.r_[estimates, targets]
labels = np.r_[np.zeros((estimates.shape[0],)), np.ones((targets.shape[0],))]
# Important: needed, since keras does not shuffle before selecting validation split
shuffle_idx = np.random.permutation(data.shape[0])
data = data[shuffle_idx]
labels = labels[shuffle_idx]
# Create and train classifier with optional stopping
def build_classifier():
classifier = keras.Sequential(
[MLP(widths=mlp_widths, **kwargs.get("mlp_kwargs", {})), keras.layers.Dense(units=1, activation="sigmoid")]
)
classifier.compile(optimizer="adam", loss="binary_crossentropy", metrics=[metric])
return classifier
if cross_validation_splits > 1:
# Create permuted indices for each class separately to ensure stratification
perm_est = np.random.permutation(len(estimates))
perm_tar = np.random.permutation(len(targets)) + len(estimates) # Offset indices
# Split indices into k folds
est_folds = np.array_split(perm_est, cross_validation_splits)
tar_folds = np.array_split(perm_tar, cross_validation_splits)
splits = []
for i in range(cross_validation_splits):
val_idx = np.concatenate([est_folds[i], tar_folds[i]])
# Create boolean mask to efficiently select training data (everything not in val)
mask = np.ones(len(data), dtype=bool)
mask[val_idx] = False
train_idx = np.where(mask)[0]
splits.append((train_idx, val_idx))
else:
# Single Split (Hold-out)
perm = np.random.permutation(len(data))
n_val = int(len(data) * validation_split)
splits = [(perm[n_val:], perm[:n_val])]
scores = []
histories = []
classifiers = []
for train_idx, val_idx in splits:
data_train, data_val = data[train_idx], data[val_idx]
labels_train, labels_val = labels[train_idx], labels[val_idx]
classifier = build_classifier()
early_stopping = keras.callbacks.EarlyStopping(
monitor=f"val_{metric}", patience=patience, restore_best_weights=True
)
# For now, we need to enable grads, since we turn them off by default
if keras.backend.backend() == "torch":
import torch
with torch.enable_grad():
history = classifier.fit(
x=data_train,
y=labels_train,
epochs=max_epochs,
batch_size=batch_size,
verbose=0,
callbacks=[early_stopping],
validation_data=(data_val, labels_val),
)
else:
history = classifier.fit(
x=data_train,
y=labels_train,
epochs=max_epochs,
batch_size=batch_size,
verbose=0,
callbacks=[early_stopping],
validation_data=(data_val, labels_val),
)
scores.append(history.history[f"val_{metric}"][-1])
if not return_metric_only:
histories.append(history.history)
classifiers.append(classifier)
scores = np.maximum(scores, 1 - np.array(scores)) # Ensure >= 0.5
mean_score = float(np.mean(scores))
if return_metric_only:
return mean_score
return {"score": mean_score, "scores": scores, "classifiers": classifiers, "histories": histories}
