|
| 1 | +from typing import Sequence, Mapping, Any |
| 2 | + |
| 3 | +import numpy as np |
| 4 | + |
| 5 | +import keras |
| 6 | + |
| 7 | +from bayesflow.utils.exceptions import ShapeError |
| 8 | +from bayesflow.networks import MLP |
| 9 | + |
| 10 | + |
| 11 | +def classifier_two_sample_test( |
| 12 | + estimates: np.ndarray, |
| 13 | + targets: np.ndarray, |
| 14 | + metric: str = "accuracy", |
| 15 | + patience: int = 5, |
| 16 | + max_epochs: int = 1000, |
| 17 | + batch_size: int = 128, |
| 18 | + return_metric_only: bool = True, |
| 19 | + validation_split: float = 0.5, |
| 20 | + standardize: bool = True, |
| 21 | + mlp_widths: Sequence = (64, 64), |
| 22 | + **kwargs, |
| 23 | +) -> float | Mapping[str, Any]: |
| 24 | + """ |
| 25 | + C2ST metric [1] between samples from two distributions computed using a neural classifier. |
| 26 | + Can be computationally expensive if called in a loop[, since it needs to train the model |
| 27 | + for each set of samples. |
| 28 | +
|
| 29 | + Note: works best for large numbers of samples and averaged across different posteriors. |
| 30 | +
|
| 31 | + [1] Lopez-Paz, D., & Oquab, M. (2016). Revisiting classifier two-sample tests. arXiv:1610.06545. |
| 32 | +
|
| 33 | + Parameters |
| 34 | + ---------- |
| 35 | + estimates : np.ndarray |
| 36 | + Array of shape (num_samples_est, num_variables) containing samples representing estimated quantities |
| 37 | + (e.g., approximate posterior samples). |
| 38 | + targets : np.ndarray |
| 39 | + Array of shape (num_samples_tar, num_variables) containing target samples |
| 40 | + (e.g., samples from a reference posterior). |
| 41 | + metric : str, optional |
| 42 | + Metric to evaluate the classifier performance. Default is "accuracy". |
| 43 | + patience : int, optional |
| 44 | + Number of epochs with no improvement after which training will be stopped. Default is 5. |
| 45 | + max_epochs : int, optional |
| 46 | + Maximum number of epochs to train the classifier. Default is 1000. |
| 47 | + batch_size : int, optional |
| 48 | + Number of samples per batch during training. Default is 64. |
| 49 | + return_metric_only : bool, optional |
| 50 | + If True, only the final validation metric is returned. Otherwise, a dictionary with the score, classifier, and |
| 51 | + full training history is returned. Default is True. |
| 52 | + validation_split : float, optional |
| 53 | + Fraction of the training data to be used as validation data. Default is 0.5. |
| 54 | + standardize : bool, optional |
| 55 | + If True, both estimates and targets will be standardized using the mean and standard deviation of estimates. |
| 56 | + Default is True. |
| 57 | + mlp_widths : Sequence[int], optional |
| 58 | + Sequence specifying the number of units in each hidden layer of the MLP classifier. Default is (256, 256). |
| 59 | + **kwargs |
| 60 | + Additional keyword arguments. Recognized keyword: |
| 61 | + mlp_kwargs : dict |
| 62 | + Dictionary of additional parameters to pass to the MLP constructor. |
| 63 | +
|
| 64 | + Returns |
| 65 | + ------- |
| 66 | + results : float or dict |
| 67 | + If return_metric_only is True, returns the final validation metric (e.g., accuracy) as a float. |
| 68 | + Otherwise, returns a dictionary with keys "score", "classifier", and "history", where "score" |
| 69 | + is the final validation metric, "classifier" is the trained Keras model, and "history" contains the |
| 70 | + full training history. |
| 71 | + """ |
| 72 | + |
| 73 | + # Error, if targets dim does not match estimates dim |
| 74 | + num_dims = estimates.shape[1] |
| 75 | + if not num_dims == targets.shape[1]: |
| 76 | + raise ShapeError( |
| 77 | + f"estimates and targets can have different number of samples (1st dim)" |
| 78 | + f"but must have the same dimensionality (2nd dim)" |
| 79 | + f"found: estimates shape {estimates.shape[1]}, targets shape {targets.shape[1]}" |
| 80 | + ) |
| 81 | + |
| 82 | + # Standardize both estimates and targets relative to estimates mean and std |
| 83 | + if standardize: |
| 84 | + estimates_mean = np.mean(estimates, axis=0) |
| 85 | + estimates_std = np.std(estimates, axis=0) |
| 86 | + estimates = (estimates - estimates_mean) / estimates_std |
| 87 | + targets = (targets - estimates_mean) / estimates_std |
| 88 | + |
| 89 | + # Create data for classification task |
| 90 | + data = np.r_[estimates, targets] |
| 91 | + labels = np.r_[np.zeros((estimates.shape[0],)), np.ones((targets.shape[0],))] |
| 92 | + |
| 93 | + # Important: needed, since keras does not shuffle before selecting validation split |
| 94 | + shuffle_idx = np.random.permutation(data.shape[0]) |
| 95 | + data = data[shuffle_idx] |
| 96 | + labels = labels[shuffle_idx] |
| 97 | + |
| 98 | + # Create and train classifier with optional stopping |
| 99 | + classifier = keras.Sequential( |
| 100 | + [MLP(widths=mlp_widths, **kwargs.get("mlp_kwargs", {})), keras.layers.Dense(1, activation="sigmoid")] |
| 101 | + ) |
| 102 | + |
| 103 | + classifier.compile(optimizer="adam", loss="binary_crossentropy", metrics=[metric]) |
| 104 | + |
| 105 | + early_stopping = keras.callbacks.EarlyStopping( |
| 106 | + monitor=f"val_{metric}", patience=patience, restore_best_weights=True |
| 107 | + ) |
| 108 | + |
| 109 | + # For now, we need to enable grads, since we turn them off by default |
| 110 | + if keras.backend.backend() == "torch": |
| 111 | + import torch |
| 112 | + |
| 113 | + with torch.enable_grad(): |
| 114 | + history = classifier.fit( |
| 115 | + x=data, |
| 116 | + y=labels, |
| 117 | + epochs=max_epochs, |
| 118 | + batch_size=batch_size, |
| 119 | + verbose=0, |
| 120 | + callbacks=[early_stopping], |
| 121 | + validation_split=validation_split, |
| 122 | + ) |
| 123 | + else: |
| 124 | + history = classifier.fit( |
| 125 | + x=data, |
| 126 | + y=labels, |
| 127 | + epochs=max_epochs, |
| 128 | + batch_size=batch_size, |
| 129 | + verbose=0, |
| 130 | + callbacks=[early_stopping], |
| 131 | + validation_split=validation_split, |
| 132 | + ) |
| 133 | + |
| 134 | + if return_metric_only: |
| 135 | + return history.history[f"val_{metric}"][-1] |
| 136 | + return {"score": history.history[f"val_{metric}"][-1], "classifier": classifier, "history": history.history} |
0 commit comments