Enhance posterior viz

brian2modelfitting/inferencer.py

@@ -161,7 +161,7 @@ class Inferencer(object):
         key corresponds to the name of the output variable as defined
         in ``model`` and value corresponds to a single dimensional
         array of recorded data traces.
-    features : list
+    features : list, optional
         List of callables that take the voltage trace and output
         summary statistics.
     method : str, optional
@@ -183,7 +183,7 @@ class Inferencer(object):
         Dictionary of state variables to be initialized with respective
         values.
     """
-    def __init__(self, dt, model, input, output, features, method=None,
+    def __init__(self, dt, model, input, output, features=None, method=None,
                  threshold=None, reset=None, refractory=False,
                  param_init=None):
         # time scale
@@ -275,13 +275,16 @@ def __init__(self, dt, model, input, output, features, method=None,
         # placeholder for the posterior
         self.posterior = None
         # observation the focus is on
-        x_o = []
-        for o in self.output:
-            o = np.array(o)
-            obs = []
-            for feature in features:
-                obs.extend(feature(o.transpose()))
-            x_o.append(obs)
+        if features:
+            x_o = []
+            for o in self.output:
+                o = np.array(o)
+                obs = []
+                for feature in features:
+                    obs.extend(feature(o.transpose()))
+                x_o.append(obs)
+        else:
+            x_o = np.vstack(self.output)
         x_o = torch.tensor(x_o, dtype=torch.float32)
         self.x_o = x_o
         self.features = features
@@ -455,11 +458,15 @@ def extract_summary_statistics(self, theta, level=1):
         x = []
         for ov in self.output_var:
             x_val = obs[ov].get_value()
-            summary_statistics = []
-            for feature in self.features:
-                summary_statistics.append(feature(x_val))
-            x.append(summary_statistics)
-        x = np.array(x, dtype=np.float32)
+            if self.features:
+                summary_statistics = []
+                for feature in self.features:
+                    summary_statistics.append(feature(x_val))
+                x.append(summary_statistics)
+                x = np.array(x, dtype=np.float32)
+            else:
+                x.append(x_val)
+                x = np.vstack(x).astype(np.float32)
         x = x.reshape((self.n_samples, -1))
         return x
 
@@ -893,7 +900,7 @@ def pairplot(self, samples=None, **kwargs):
             Samples used to build the pairplot.
         **kwargs : dict, optional
             Additional keyword arguments for the
-            ``sbi.analysis.plot.pairplot`` function.
+            ``sbi.analysis.pairplot`` function.
 
         Returns
         -------
@@ -911,6 +918,82 @@ def pairplot(self, samples=None, **kwargs):
         fig, axes = sbi.analysis.pairplot(s, **kwargs)
         return fig, axes
 
+    def conditional_pairplot(self, condition, limits, density=None, **kwargs):
+        """Plot conditional distribution given all other parameters.
+
+        Check ``sbi.analysis.plot.conditional_pairplot`` for more
+        details.
+
+        Parameters
+        ----------
+        condition : torch.tensor
+            Condition that all but the one/two regarded parameters are
+            fixed to.
+        limits : list or torch.tensor
+            Limits in between which each parameter will be evaulated.
+        density : sbi.inference.NeuralPosterior, optional
+            Posterior probability density.
+        **kwargs : dict, optional
+            Additional keyword arguments for the
+            ``sbi.analysis.conditional_pairplot`` function.
+
+        Returns
+        -------
+        tuple
+            Figure and axis of conditional pairplot.
+        """
+        if density is not None:
+            d = density
+        else:
+            try:
+                d = self.posterior
+            except AttributeError as e:
+                print(e, '\nDensity is not available.')
+                raise
+        fig, axes = sbi.analysis.conditional_pairplot(density=d,
+                                                      condition=condition,
+                                                      limits=limits,
+                                                      *kwargs)
+        return fig, axes
+
+    def conditional_corrcoeff(self, condition, limits, density=None, **kwargs):
+        """Plot conditional distribution given all other parameters.
+
+        Check ``sbi.analysis.conditional_density.conditional_corrcoeff``
+        for more details.
+
+        Parameters
+        ----------
+        condition : torch.tensor
+            Condition that all but the one/two regarded parameters are
+            fixed to.
+        limits : list or torch.tensor
+            Limits in between which each parameter will be evaulated.
+        density : sbi.inference.NeuralPosterior, optional
+            Posterior probability density.
+        **kwargs : dict, optional
+            Additional keyword arguments for the
+            ``sbi.analysis.conditional_corrcoeff`` function.
+
+        Returns
+        -------
+        torch.tensor
+            Average conditional correlation matrix.
+        """
+        if density is not None:
+            d = density
+        else:
+            try:
+                d = self.posterior
+            except AttributeError as e:
+                print(e, '\nDensity is not available.')
+                raise
+        cond_coeff = sbi.analysis.conditional_corrcoeff(density=d,
+                                                        condition=condition,
+                                                        limits=limits,
+                                                        *kwargs)
+        return cond_coeff
+
     def generate_traces(self, posterior=None, output_var=None, param_init=None,
                         level=0):
         """Generates traces for a single drawn sample from the trained