from collections.abc import Sequence, Mapping, Callable
import numpy as np
from bayesflow.utils import logging
from ...utils.dict_utils import dicts_to_arrays, compute_test_quantities
[docs]
def root_mean_squared_error(
estimates: Mapping[str, np.ndarray] | np.ndarray,
targets: Mapping[str, np.ndarray] | np.ndarray,
variable_keys: Sequence[str] = None,
variable_names: Sequence[str] = None,
test_quantities: dict[str, Callable] = None,
normalize: str | None = "prior",
aggregation: Callable = np.median,
) -> dict[str, any]:
"""
Computes the (Normalized) Root Mean Squared Error (RMSE/NRMSE) for the given posterior and prior samples.
The values of the default normalization (`prior`) should be interpreted as 0 indicating the most informative
(posterior is a point mass at ground truth) and 1 indicating non-informative (posterior equals prior) results.
Parameters
----------
estimates : np.ndarray or dict[str, np.ndarray]
Posterior samples, either as a NumPy array of shape (num_datasets, num_draws_post, num_variables)
or as a dictionary mapping variable names to arrays.
Comprises `num_draws_post` random draws from the posterior distribution
for each data set from `num_datasets`.
targets : np.ndarray or dict[str, np.ndarray]
Prior samples, either as a NumPy array of shape (num_datasets, num_variables)
or as a dictionary mapping variable names to arrays.
Comprises `num_datasets` ground truths.
variable_keys : Sequence[str], optional (default = None)
Select keys from the dictionaries provided in estimates and targets.
By default, select all keys.
variable_names : Sequence[str], optional (default = None)
Optional variable names to show in the output.
test_quantities : dict or None, optional, default: None
A dict that maps plot titles to functions that compute
test quantities based on estimate/target draws.
The dict keys are automatically added to ``variable_keys``
and ``variable_names``.
Test quantity functions are expected to accept a dict of draws with
shape ``(batch_size, ...)`` as the first (typically only)
positional argument and return an NumPy array of shape
``(batch_size,)``.
The functions do not have to deal with an additional
sample dimension, as appropriate reshaping is done internally.
normalize : str or None, optional (default = "prior")
Whether to normalize the RMSE using statistics of the prior samples.
Possible options are ("mean", "range", "median", "iqr", "std", "prior", None)
aggregation : Callable, optional (default = np.median)
Function to aggregate the RMSE across draws. Typically `np.mean` or `np.median`.
Notes
-----
Aggregation is performed after computing the RMSE for each posterior draw, instead of first aggregating
the posterior draws and then computing the RMSE between aggregates and ground truths.
Returns
-------
result : dict
Dictionary containing:
- "values" : np.ndarray
The aggregated (N)RMSE for each variable.
- "metric_name" : str
The name of the metric ("RMSE" or "NRMSE").
- "variable_names" : str
The (inferred) variable names.
"""
if normalize:
logging.warning(
"Using new default normalize='prior' for a more dynamic range. "
"To reproduce previous behavior, set normalize='range'."
)
# Optionally, compute and prepend test quantities from draws
if test_quantities is not None:
updated_data = compute_test_quantities(
targets=targets,
estimates=estimates,
variable_keys=variable_keys,
variable_names=variable_names,
test_quantities=test_quantities,
)
variable_names = updated_data["variable_names"]
variable_keys = updated_data["variable_keys"]
estimates = updated_data["estimates"]
targets = updated_data["targets"]
samples = dicts_to_arrays(
estimates=estimates,
targets=targets,
variable_keys=variable_keys,
variable_names=variable_names,
)
err = samples["estimates"] - samples["targets"][:, None, :]
rmse = np.sqrt(np.mean(err**2, axis=1))
targets = samples["targets"]
match normalize:
case None | False:
normalizer = np.array(1.0)
case "mean":
normalizer = np.mean(targets, axis=0)
case "median":
normalizer = np.median(targets, axis=0)
case "range":
normalizer = targets.max(axis=0) - targets.min(axis=0)
case "std":
normalizer = np.std(targets, axis=0, ddof=0)
case "iqr":
q75 = np.percentile(targets, 75, axis=0)
q25 = np.percentile(targets, 25, axis=0)
normalizer = q75 - q25
case "prior":
N, S, _ = samples["estimates"].shape
# bootstrap prior-only predictions from empirical prior samples in targets
idx = np.random.randint(0, N, size=(N, S))
prior_bootstrap = targets[idx]
prior_err = prior_bootstrap - targets[:, None, :]
prior_rmse = np.sqrt(np.mean(prior_err**2, axis=1))
normalizer = aggregation(prior_rmse, axis=0)
case _:
raise ValueError(f"Unknown normalization mode: {normalize}")
metric_name = "NRMSE" if normalize else "RMSE"
rmse = rmse / normalizer
rmse = aggregation(rmse, axis=0)
variable_names = samples["estimates"].variable_names
return {"values": rmse, "metric_name": metric_name, "variable_names": variable_names}
