Repository: ICB-DCM/pyABC

Issue:
Log distance weights
Until the near future, the distance weights will not be stored in the history due to legacy reasons. Thus, it may be good to add a callback function which is called after each update method.

Sidenote: Always possible by creating a SubClass, or, to some degree, using the `logging.getLogger("Distance")` logger.

Comments:

Patch:
diff --git a/pyabc/acceptor/acceptor.py b/pyabc/acceptor/acceptor.py
index cc8ca36..3c6acac 100644
--- a/pyabc/acceptor/acceptor.py
+++ b/pyabc/acceptor/acceptor.py
@@ -19,7 +19,7 @@ import pandas as pd
 from typing import Callable
 import logging
 
-from ..distance import Distance, SCALE_LIN
+from ..distance import Distance, SCALE_LIN, StochasticKernel
 from ..epsilon import Epsilon
 from ..parameters import Parameter
 from .pdf_norm import pdf_norm_max_found
@@ -354,7 +354,8 @@ class StochasticAcceptor(Acceptor):
             usually only for testing purposes.
         log_file: str, optional
             A log file for storing data of the acceptor that are currently not
-            saved in the database. The data are saved in json format.
+            saved in the database. The data are saved in json format and can
+            be retrieved via `pyabc.storage.load_dict_from_json`.
         """
         super().__init__()
 
@@ -376,7 +377,7 @@ class StochasticAcceptor(Acceptor):
             self,
             t: int,
             get_weighted_distances: Callable[[], pd.DataFrame],
-            distance_function: Distance,
+            distance_function: StochasticKernel,
             x_0: dict):
         """
         Initialize temperature and maximum pdf.
@@ -430,7 +431,9 @@ class StochasticAcceptor(Acceptor):
             kernel_scale=self.kernel_scale,  # TODO Refactor
         )
 
-    def __call__(self, distance_function, eps, x, x_0, t, par):
+    def __call__(self,
+                 distance_function: StochasticKernel,
+                 eps, x, x_0, t, par):
         # rename
         kernel = distance_function
 
@@ -438,15 +441,15 @@ class StochasticAcceptor(Acceptor):
         temp = eps(t)
 
         # compute probability density
-        pd = kernel(x, x_0, t, par)
+        density = kernel(x, x_0, t, par)
 
         pdf_norm = self.pdf_norms[t]
 
         # compute acceptance probability
         if kernel.ret_scale == SCALE_LIN:
-            acc_prob = (pd / pdf_norm) ** (1 / temp)
+            acc_prob = (density / pdf_norm) ** (1 / temp)
         else:  # kernel.ret_scale == SCALE_LOG
-            acc_prob = np.exp((pd - pdf_norm) * (1 / temp))
+            acc_prob = np.exp((density - pdf_norm) * (1 / temp))
 
         # accept
         threshold = np.random.uniform(low=0, high=1)
@@ -464,10 +467,10 @@ class StochasticAcceptor(Acceptor):
             weight = 1.0
 
         # check pdf max ok
-        if pdf_norm < pd:
+        if pdf_norm < density:
             logger.debug(
-                f"Encountered pd={pd:.4e} > c={pdf_norm:.4e}, "
+                f"Encountered density={density:.4e} > c={pdf_norm:.4e}, "
                 f"thus weight={weight:.4e}.")
 
         # return unscaled density value and the acceptance flag
-        return AcceptorResult(pd, accept, weight)
+        return AcceptorResult(density, accept, weight)
diff --git a/pyabc/distance/distance.py b/pyabc/distance/distance.py
index 992a2f6..2471e61 100644
--- a/pyabc/distance/distance.py
+++ b/pyabc/distance/distance.py
@@ -6,6 +6,7 @@ import logging
 from ..sampler import Sampler
 from .scale import standard_deviation, span
 from .base import Distance, to_distance
+from ..storage import save_dict_to_json
 
 
 logger = logging.getLogger("Distance")
@@ -141,32 +142,37 @@ class AdaptivePNormDistance(PNormDistance):
     Parameters
     ----------
 
-    p: float, optional (default = 2)
+    p:
         p for p-norm. Required p >= 1, p = np.inf allowed (infinity-norm).
-    initial_weights: dict, optional
+        Default: p=2.
+    initial_weights:
         Weights to be used in the initial iteration. Dictionary with
         observables as keys and weights as values.
-    factors: dict, optional
+    factors:
         As in PNormDistance.
-    adaptive: bool, optional (default = True)
+    adaptive:
         True: Adapt distance after each iteration.
         False: Adapt distance only once at the beginning in initialize().
         This corresponds to a pre-calibration.
-    scale_function: Callable, optional (default = standard_deviation)
+    scale_function:
         (data: list, x_0: float) -> scale: float. Computes the scale (i.e.
         inverse weight s = 1 / w) for a given summary statistic. Here, data
         denotes the list of simulated summary statistics, and x_0 the observed
         summary statistic. Implemented are absolute_median_deviation,
         standard_deviation (default), centered_absolute_median_deviation,
         centered_standard_deviation.
-    normalize_weights: bool, optional (default = True)
+    normalize_weights:
         Whether to normalize the weights to have mean 1. This just possibly
         smoothes the decrease of epsilon and might aid numeric stability, but
         is not strictly necessary.
-    max_weight_ratio: float, optional (default = None)
+    max_weight_ratio:
         If not None, large weights will be bounded by the ratio times the
         smallest non-zero absolute weight. In practice usually not necessary,
         it is theoretically required to ensure convergence.
+    log_file:
+        A log file to store weights for each time point in. Weights are
+        currently not stored in the database. The data are saved in json
+        format and can be retrieved via `pyabc.storage.load_dict_from_json`.
 
 
     .. [#prangle] Prangle, Dennis. "Adapting the ABC Distance Function".
@@ -178,9 +184,10 @@ class AdaptivePNormDistance(PNormDistance):
                  initial_weights: dict = None,
                  factors: dict = None,
                  adaptive: bool = True,
-                 scale_function=None,
+                 scale_function: Callable = None,
                  normalize_weights: bool = True,
-                 max_weight_ratio: float = None):
+                 max_weight_ratio: float = None,
+                 log_file: str = None):
         # call p-norm constructor
         super().__init__(p=p, weights=None, factors=factors)
 
@@ -194,6 +201,7 @@ class AdaptivePNormDistance(PNormDistance):
 
         self.normalize_weights = normalize_weights
         self.max_weight_ratio = max_weight_ratio
+        self.log_file = log_file
 
         self.x_0 = None
 
@@ -294,7 +302,7 @@ class AdaptivePNormDistance(PNormDistance):
         self.weights[t] = w
 
         # logging
-        logger.debug(f"updated weights[{t}] = {self.weights[t]}")
+        self.log(t)
 
     def _normalize_weights(self, w):
         """
@@ -346,6 +354,12 @@ class AdaptivePNormDistance(PNormDistance):
                 "normalize_weights": self.normalize_weights,
                 "max_weight_ratio": self.max_weight_ratio}
 
+    def log(self, t: int) -> None:
+        logger.debug(f"updated weights[{t}] = {self.weights[t]}")
+
+        if self.log_file:
+            save_dict_to_json(self.weights, self.log_file)
+
 
 class AggregatedDistance(Distance):
     """
@@ -383,6 +397,8 @@ class AggregatedDistance(Distance):
             can remain static while weights adapt over time, allowing for
             greater flexibility.
         """
+        super().__init__()
+
         if not isinstance(distances, list):
             distances = [distances]
         self.distances = [to_distance(distance) for distance in distances]
@@ -440,7 +456,7 @@ class AggregatedDistance(Distance):
         self.format_weights_and_factors(t)
         weights = AggregatedDistance.get_for_t_or_latest(self.weights, t)
         factors = AggregatedDistance.get_for_t_or_latest(self.factors, t)
-        return np.dot(weights * factors, values)
+        return float(np.dot(weights * factors, values))
 
     def get_config(self) -> dict:
         """
@@ -495,23 +511,26 @@ class AdaptiveAggregatedDistance(AggregatedDistance):
 
     Parameters
     ----------
-
-    distances: List[Distance]
+    distances:
         As in AggregatedDistance.
-    initial_weights: list, optional
+    initial_weights:
         Weights to be used in the initial iteration. List with
         a weight for each distance function.
-    factors: Union[List, dict]
+    factors:
         As in AggregatedDistance.
-    adaptive: bool, optional (default = True)
+    adaptive:
         True: Adapt weights after each iteration.
         False: Adapt weights only once at the beginning in initialize().
         This corresponds to a pre-calibration.
-    scale_function: Callable, optional (default = scale_span)
+    scale_function:
         Function that takes a list of floats, namely the values obtained
         by applying one of the distances passed to a set of samples,
         and returns a single float, namely the weight to apply to this
-        distance function.
+        distance function. Default: scale_span.
+    log_file:
+        A log file to store weights for each time point in. Weights are
+        currently not stored in the database. The data are saved in json
+        format and can be retrieved via `pyabc.storage.load_dict_from_json`.
     """
 
     def __init__(
@@ -520,7 +539,8 @@ class AdaptiveAggregatedDistance(AggregatedDistance):
             initial_weights: List = None,
             factors: Union[List, dict] = None,
             adaptive: bool = True,
-            scale_function: Callable = None):
+            scale_function: Callable = None,
+            log_file: str = None):
         super().__init__(distances=distances)
         self.initial_weights = initial_weights
         self.factors = factors
@@ -529,6 +549,7 @@ class AdaptiveAggregatedDistance(AggregatedDistance):
         if scale_function is None:
             scale_function = span
         self.scale_function = scale_function
+        self.log_file = log_file
 
     def initialize(self,
                    t: int,
@@ -599,8 +620,14 @@ class AdaptiveAggregatedDistance(AggregatedDistance):
         self.weights[t] = np.array(w)
 
         # logging
+        self.log(t)
+
+    def log(self, t: int) -> None:
         logger.debug(f"updated weights[{t}] = {self.weights[t]}")
 
+        if self.log_file:
+            save_dict_to_json(self.weights, self.log_file)
+
 
 class DistanceWithMeasureList(Distance):
     """


Test patch:
diff --git a/test/test_distance_function.py b/test/test_distance_function.py
index e40c2fc..ef95c37 100644
--- a/test/test_distance_function.py
+++ b/test/test_distance_function.py
@@ -1,5 +1,8 @@
 import numpy as np
 import scipy as sp
+import scipy.stats
+import tempfile
+
 from pyabc.distance import (
     PercentileDistance,
     MinMaxDistance,
@@ -27,6 +30,7 @@ from pyabc.distance import (
     mean,
     median,
     SCALE_LIN,)
+from pyabc.storage import load_dict_from_json
 
 
 class MockABC:
@@ -405,3 +409,31 @@ def test_negativebinomialkernel():
     expected = np.sum(sp.stats.nbinom.logpmf(
         k=[4, 5, 8], n=[7, 7, 7], p=[0.9, 0.8, 0.7]))
     assert np.isclose(ret, expected)
+
+
+def test_store_weights():
+    """Test whether storing distance weights works."""
+    abc = MockABC([{'s1': -1, 's2': -1, 's3': -1},
+                   {'s1': -1, 's2': 0, 's3': 1}])
+    x_0 = {'s1': 0, 's2': 0, 's3': 1}
+
+    weights_file = tempfile.mkstemp(suffix=".json")[1]
+    print(weights_file)
+
+    def distance0(x, x_0):
+        return abs(x['s0'] - x_0['s0'])
+
+    def distance1(x, x_0):
+        return np.sqrt((x['s1'] - x_0['s1'])**2)
+
+    for distance in [AdaptivePNormDistance(log_file=weights_file),
+                     AdaptiveAggregatedDistance(
+                         [distance0, distance1], log_file=weights_file)]:
+        distance = AdaptivePNormDistance(log_file=weights_file)
+        distance.initialize(0, abc.sample_from_prior, x_0=x_0)
+        distance.update(1, abc.sample_from_prior)
+        distance.update(2, abc.sample_from_prior)
+
+        weights = load_dict_from_json(weights_file)
+        assert set(weights.keys()) == {0, 1, 2}
+        assert weights == distance.weights