Merge pull request #147 from han-ol/Development

Bayes Estimators: Loss functions with flexible signature

Author: stefanradev93

README.md

@@ -105,10 +105,10 @@ the model-amortizer combination on unseen simulations:
 # Generate 500 new simulated data sets
 new_sims = trainer.configurator(generative_model(500))
 
-# Obtain 100 posteriors draws per data set instantly
+# Obtain 100 posterior draws per data set instantly
 posterior_draws = amortized_posterior.sample(new_sims, n_samples=100)
 
-# Diagnoze calibration
+# Diagnose calibration
 fig = bf.diagnostics.plot_sbc_histograms(posterior_draws, new_sims['parameters'])
 ```
 
@@ -302,7 +302,7 @@ This project is currently managed by researchers from Rensselaer Polytechnic Ins
 You can cite BayesFlow along the lines of:
 
 - We approximated the posterior with neural posterior estimation and learned summary statistics (NPE; Radev et al., 2020), as implemented in the BayesFlow software for amortized Bayesian workflows (Radev et al., 2023a).
-- We approximated the likelihood with neural likelihood estimation (NLE; Papamakarios et al., 2019) without hand-cafted summary statistics, as implemented in the BayesFlow software for amortized Bayesian workflows (Radev et al., 2023b).
+- We approximated the likelihood with neural likelihood estimation (NLE; Papamakarios et al., 2019) without hand-crafted summary statistics, as implemented in the BayesFlow software for amortized Bayesian workflows (Radev et al., 2023b).
 - We performed simultaneous posterior and likelihood estimation with jointly amortized neural approximation (JANA; Radev et al., 2023a), as implemented in the BayesFlow software for amortized Bayesian workflows (Radev et al., 2023b).
 
 1. Radev, S. T., Schmitt, M., Schumacher, L., Elsemüller, L., Pratz, V., Schälte, Y., Köthe, U., & Bürkner, P.-C. (2023a). BayesFlow: Amortized Bayesian workflows with neural networks. *The Journal of Open Source Software, 8(89)*, 5702.([arXiv](https://arxiv.org/abs/2306.16015))([JOSS](https://joss.theoj.org/papers/10.21105/joss.05702))

bayesflow/amortizers.py

@@ -31,7 +31,7 @@
 from bayesflow.default_settings import DEFAULT_KEYS
 from bayesflow.exceptions import ConfigurationError, SummaryStatsError
 from bayesflow.helper_functions import check_tensor_sanity
-from bayesflow.losses import log_loss, mmd_summary_space
+from bayesflow.losses import log_loss, mmd_summary_space, norm_diff
 from bayesflow.networks import EvidentialNetwork
 
 
@@ -1337,7 +1337,7 @@ def compute_loss(self, input_dict, **kwargs):
         """
 
         net_out = self(input_dict, **kwargs)
-        loss = tf.reduce_mean(self.loss_fn(net_out - input_dict[DEFAULT_KEYS["parameters"]]))
+        loss = tf.reduce_mean(self.loss_fn(net_out, input_dict[DEFAULT_KEYS["parameters"]]))
         return loss
 
     def _compute_summary_condition(self, summary_conditions, direct_conditions, **kwargs):
@@ -1366,4 +1366,4 @@ def _determine_loss(self, loss_fun, norm_ord):
         # In case of user-provided loss, override norm order
         if loss_fun is not None:
             return loss_fun
-        return partial(tf.norm, ord=norm_ord, axis=-1)
+        return partial(norm_diff, ord=norm_ord, axis=-1)

bayesflow/losses.py

@@ -184,3 +184,20 @@ def log_loss(model_indices, preds, evidential=False, label_smoothing=0.01):
     # Actual loss + regularization (if given)
     loss = -tf.reduce_mean(tf.reduce_sum(model_indices * tf.math.log(preds), axis=1))
     return loss
+
+
+def norm_diff(tensor_a, tensor_b, axis=None, ord='euclidean'):
+    """
+    Wrapper around tf.norm that computes the norm of the difference between two tensors along the specified axis.
+
+    Parameters
+    ----------
+    tensor_a : A Tensor.
+    tensor_b : A Tensor. Must be the same shape as tensor_a.
+    axis     : Any or None
+        Axis along which to compute the norm of the difference. Default is None.
+    ord      : int or str
+        Order of the norm. Supports 'euclidean' and other norms supported by tf.norm. Default is 'euclidean'.
+    """
+    difference = tensor_a - tensor_b
+    return tf.norm(difference, ord=ord, axis=axis)