Merge pull request #75 from IMAGINE-Consortium/pipeline_tests

Includes tools for testing the pipeline and checking likelihood convergence

imagine/likelihoods/ensemble_likelihood.py

@@ -50,10 +50,10 @@ class EnsembleLikelihood(Likelihood):
         covariance matrix from the simulations.
     """
     def __init__(self, measurement_dict, covariance_dict=None, mask_dict=None,
-                 cov_func=None, use_trace_approximation=False):
+                 cov_func=None, use_trace_approximation=False, **kwargs):
 
         super().__init__(measurement_dict, covariance_dict=covariance_dict,
-                         mask_dict=mask_dict)
+                         mask_dict=mask_dict, **kwargs)
 
         # Requires covariaces to be present when using this type of Likelihood
         assert self._covariance_dict is not None
@@ -125,10 +125,10 @@ class EnsembleLikelihoodDiagonal(Likelihood):
         Masks which will be applied to the Measurements, Covariances and
         Simulations, before computing the likelihood
     """
-    def __init__(self, measurement_dict, covariance_dict=None, mask_dict=None):
+    def __init__(self, measurement_dict, covariance_dict=None, mask_dict=None, **kwargs):
 
         super().__init__(measurement_dict, covariance_dict=covariance_dict,
-                         mask_dict=mask_dict)
+                         mask_dict=mask_dict, **kwargs)
 
         # Requires covariaces to be present when using this type of Likelihood
         assert self._covariance_dict is not None

imagine/likelihoods/likelihood.py

@@ -21,6 +21,9 @@
 # Built-in imports
 import abc
 
+# Package imports
+import numpy as np
+
 # IMAGINE imports
 from imagine.observables.observable_dict import (
     Measurements, Covariances, Masks)
@@ -50,27 +53,56 @@ class Likelihood(BaseClass, metaclass=abc.ABCMeta):
     mask_dict : imagine.observables.observable_dict.Masks
         A :py:obj:`Masks <imagine.observables.observable_dict.Masks>` dictionary
         which should be applied to the measured and simulated data.
+    compute_dispersion : bool
+        If True, calling the Likelihood object will return the likelihood value
+        and the dispersion estimated by bootstrapping the simulations object
+        and computing the sample standard deviation.
+        If False (default), only the likelihood value is returned.
+    n_bootstrap : int
+        Number of resamples used in the bootstrapping of the simulations if
+        compute_dispersion is set to `True`.
     """
-
-    def __init__(self, measurement_dict, covariance_dict=None, mask_dict=None):
+    def __init__(self, measurement_dict, covariance_dict=None, mask_dict=None,
+                 compute_dispersion=False, n_bootstrap=150):
         # Call super constructor
         super().__init__()
 
+        self._check_units(measurement_dict, covariance_dict)
         self.mask_dict = mask_dict
         self.measurement_dict = measurement_dict
         if covariance_dict is None:
             covariance_dict = measurement_dict.cov
         self.covariance_dict = covariance_dict
+        self.compute_dispersion = compute_dispersion
+        self.n_bootstrap = n_bootstrap
 
     def __call__(self, observable_dict, **kwargs):
         if self.mask_dict is not None:
             observable_dict = self.mask_dict(observable_dict)
-        return(self.call(observable_dict, **kwargs))
+
+        likelihood = self.call(observable_dict, **kwargs)
+
+        if not self.compute_dispersion:
+            return likelihood
+        else:
+            bootstrap_sample = [self._bootstrapped_likelihood(observable_dict,
+                                                              **kwargs)
+                                for _ in range(self.n_bootstrap)]
+            dispersion = np.std(bootstrap_sample)
+            return likelihood, dispersion
+
+    def _bootstrapped_likelihood(self, simulations, **kwargs):
+        # Gets ensemble size from first entry in the ObservableDict
+        size, _ = simulations[list(simulations.keys())[0]].shape
+        # Resamples with replacement
+        idx = np.random.randint(0, size, size)
+        sims_new = simulations.sub_sim(idx)
+        return self.call(sims_new, **kwargs)
 
     @property
     def mask_dict(self):
         """
-        :py:obj:`Masks <imagine.observables.observable_dict.Masks>` dictionary associated with 
+        :py:obj:`Masks <imagine.observables.observable_dict.Masks>` dictionary associated with
         this object
         """
         return self._mask_dict
@@ -84,9 +116,9 @@ def mask_dict(self, mask_dict):
     @property
     def measurement_dict(self):
         """
-        :py:obj:`Measurements <imagine.observables.observable_dict.Measurements>` dictionary associated with 
+        :py:obj:`Measurements <imagine.observables.observable_dict.Measurements>` dictionary associated with
         this object
-        
+
         NB If a mask is used, only the masked version is stored
         """
         return self._measurement_dict
@@ -101,9 +133,9 @@ def measurement_dict(self, measurement_dict):
     @property
     def covariance_dict(self):
         """
-        :py:obj:`Covariances <imagine.observables.observable_dict.Covariances>` dictionary associated with 
+        :py:obj:`Covariances <imagine.observables.observable_dict.Covariances>` dictionary associated with
         this object
-        
+
         NB If a mask is used, only the masked version is stored
         """
         return self._covariance_dict
@@ -125,3 +157,16 @@ def call(self, observable_dict):
         variables
         """
         raise NotImplementedError
+
+    @staticmethod
+    def _check_units(measurements, covariances):
+        """
+        Makes sure that measurements and covariances units are compatible
+        """
+        if covariances is None:
+            return
+        for k in measurements:
+            if measurements[k].unit is None:
+                assert covariances[k].unit is None
+            else:
+                assert (measurements[k].unit)**2 == covariances[k].unit

imagine/observables/dataset.py

@@ -67,7 +67,7 @@ def key(self):
     @property
     def var(self):
         if (self.cov is None) and (self._error is not None):
-            self._var = np.ones_like(self._data) * self._error**2
+            self._var = np.ones(self._data.shape) * self._error**2
         return self._var
 
 

imagine/observables/observable_dict.py

@@ -43,7 +43,7 @@
 # IMAGINE imports
 from imagine.observables.dataset import Dataset, HEALPixDataset
 from imagine.observables.observable import Observable
-from imagine.tools import BaseClass, mask_cov, mask_obs, req_attr
+from imagine.tools import BaseClass, mask_cov, mask_obs, req_attr, oas_cov
 import imagine.tools.visualization as visu
 
 # All declaration
@@ -283,7 +283,6 @@ class Simulations(ObservableDict):
     See `imagine.observables.observable_dict` module documentation for
     further details.
     """
-
     def append(self, *args, **kwargs):
         log.debug('@ observable_dict::Simulations::append')
         name, data, _, otype, coords = super().append(*args, **kwargs)
@@ -303,6 +302,56 @@ def append(self, *args, **kwargs):
                 raise TypeError('unsupported data type')
 
 
+    def estimate_covariances(self, cov_est=oas_cov):
+        """
+        Produces a Covariances object based on the present Simulations
+
+        Parameters
+        ----------
+        cov_est : func
+            A function that computes the covariance given a
+            matrix of data
+
+        Returns
+        -------
+        covs : imagine.observables.Covariances
+            IMAGINE Covariances object
+        """
+        covs = Covariances()
+
+        for k in self:
+            obs = self[k]
+            d = cov_est(obs.data)*(obs.unit)**2
+            covs.append(name=k, cov_data=d, otype=obs.otype)
+
+        return covs
+
+
+    def sub_sim(self, indices):
+        """
+        Creates a new :py:obj:`Simulations <imagine.observables.Simulations>` object
+        based on a subset of the ensemble of a larger :py:obj:`Simulations <imagine.observables.Simulations>`.
+
+        Parameters
+        ----------
+        indices
+            A tuple of indices numbers, a slice object or a boolean array which
+            will be used to select the data for the sub-simulation
+
+        Returns
+        -------
+        sims_subset : imagine.observables.Simulations
+            The selected sub-simulation
+        """
+        sims_subset = Simulations()
+
+        for k in self:
+            sims_subset.append(name=k,
+                               data=self[k].global_data[indices,:],
+                               otype=self[k].otype)
+        return sims_subset
+
+
 class Covariances(ObservableDict):
     """
     Stores observational covariances