Feat/plot cleanup

notebooks/1B - An MNIST example.ipynb

@@ -9,6 +9,8 @@
    "source": [
     "## Example B - MNIST & CNN\n",
     "\n",
+    "**(works only with Swyft 0.4.4; to be updated soon)**\n",
+    "\n",
     "Authors: Noemi Anau Montel, James Alvey, Christoph Weniger\n",
     "\n",
     "Last update: 27 April 2023"
@@ -2084,7 +2086,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.9.13"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {

swyft/lightning/data.py

@@ -30,8 +30,7 @@ class SwyftDataModule(pl.LightningDataModule):
 
     Args:
         data: Simulation data
-        lenghts: List of number of samples used for [training, validation, testing].
-        fractions: Fraction of samples used for [training, validation, testing].
+        val_fraction: Fraction of data used for validation.
         batch_size: Minibatch size.
         num_workers: Number of workers for dataloader.
         shuffle: Shuffle training data.
@@ -43,18 +42,26 @@ class SwyftDataModule(pl.LightningDataModule):
     def __init__(
         self,
         data,
-        lengths: Union[Sequence[int], None] = None,
-        fractions: Union[Sequence[float], None] = None,
+        #lengths: Union[Sequence[int], None] = None,
+        #fractions: Union[Sequence[float], None] = None,
+        val_fraction: float = 0.2,
         batch_size: int = 32,
         num_workers: int = 0,
         shuffle: bool = False,
         on_after_load_sample: Optional[callable] = None,
     ):
         super().__init__()
         self.data = data
+        # TODO: Clean up codes
+        lengths = None
+        fractions = [1-val_fraction, val_fraction]
         if lengths is not None and fractions is None:
+            assert len(lengths) == 2,  "SwyftDataModule only provides training and validation data."
+            lengths = [lengths[0], lenghts[1], 0]
             self.lengths = lengths
         elif lengths is None and fractions is not None:
+            assert len(fractions) == 2, "SwyftDataModule only provides training and validation data."
+            fractions = [fractions[0], fractions[1], 0]
             self.lengths = self._get_lengths(fractions, len(data))
         else:
             raise ValueError("Either lenghts or fraction must be set, but not both.")
@@ -114,13 +121,14 @@ def val_dataloader(self):
         return dataloader
 
     def test_dataloader(self):
-        dataloader = torch.utils.data.DataLoader(
-            self.dataset_test,
-            batch_size=self.batch_size,
-            shuffle=False,
-            num_workers=self.num_workers,
-        )
-        return dataloader
+        return
+#        dataloader = torch.utils.data.DataLoader(
+#            self.dataset_test,
+#            batch_size=self.batch_size,
+#            shuffle=False,
+#            num_workers=self.num_workers,
+#        )
+#        return dataloader
 
 
 class SamplesDataset(torch.utils.data.Dataset):

swyft/lightning/utils.py

@@ -31,7 +31,6 @@
 
 from swyft.lightning.data import *
 import swyft.lightning.simulator
-from swyft.plot.mass import get_empirical_z_score
 
 import scipy
 from scipy.ndimage import gaussian_filter1d, gaussian_filter
@@ -54,28 +53,30 @@ class SwyftParameterError(Exception):
 ############################
 
 
-def _pdf_from_weighted_samples(v, w, bins=50, smooth=0, v_aux=None):
+def _pdf_from_weighted_samples(v, w, bins=50, smooth=0, smooth_prior=False):
     """Take weighted samples and turn them into a pdf on a grid.
 
     Args:
         bins
     """
     ndim = v.shape[-1]
-    if v_aux is None:
+    if not smooth_prior:
         return _weighted_smoothed_histogramdd(v, w, bins=bins, smooth=smooth)
     else:
-        h, xy = _weighted_smoothed_histogramdd(v_aux, None, bins=bins, smooth=smooth)
+        h, xy = _weighted_smoothed_histogramdd(v, w * 0 + 1, bins=bins, smooth=smooth)
         if ndim == 2:
             X, Y = np.meshgrid(xy[:, 0], xy[:, 1])
             n = len(xy)
             out = scipy.interpolate.griddata(
                 v, w, (X.flatten(), Y.flatten()), method="cubic", fill_value=0.0
             ).reshape(n, n)
+            out = out * h.numpy()
             return out, xy
         elif ndim == 1:
             out = scipy.interpolate.griddata(
                 v[:, 0], w, xy[:, 0], method="cubic", fill_value=0.0
             )
+            out = out * h.numpy()
             return out, xy
         else:
             raise KeyError("Not supported")
@@ -113,6 +114,7 @@ def get_pdf(
     aux=None,
     bins: int = 50,
     smooth: float = 0.0,
+    smooth_prior=False,
 ):
     """Generate binned PDF based on input
 
@@ -121,6 +123,7 @@ def get_pdf(
         params: Parameter names
         bins: Number of bins
         smooth: Apply Gaussian smoothing
+        smooth_prior: Smooth prior instead of posterior
 
     Returns:
         np.array, np.array: Returns densities and parameter grid.
@@ -130,7 +133,9 @@ def get_pdf(
         z_aux, _ = get_weighted_samples(aux, params)
     else:
         z_aux = None
-    return _pdf_from_weighted_samples(z, w, bins=bins, smooth=smooth, v_aux=z_aux)
+    return _pdf_from_weighted_samples(
+        z, w, bins=bins, smooth=smooth, smooth_prior=smooth_prior
+    )
 
 
 def _get_weights(logratios, normalize: bool = False):
@@ -528,4 +533,5 @@ class OnFitEndLoadBestModel:
     def on_fit_end(self):
         self.best_model_path = self.trainer.checkpoint_callback.best_model_path
         checkpoint = torch.load(self.best_model_path)
+        print("Reloading best model:", self.best_model_path)
         self.load_state_dict(checkpoint["state_dict"])

swyft/plot/__init__.py

@@ -1,19 +1,9 @@
-# from swyft.plot.constraint import diagonal_constraint, lower_constraint
-# from swyft.plot.histogram import corner, hist1d
-# from swyft.plot.mass import empirical_z_score_corner, plot_empirical_z_score
-# from swyft.plot.violin import violin
-from swyft.plot.plot2 import plot_1d, plot_2d, corner, plot_zz, plot_pp
+from swyft.plot.plot import (
+    plot_posterior,
+    plot_corner,
+    plot_zz,
+    plot_pp,
+    plot_pair,
+)
 
-__all__ = [
-    "plot_1d",
-    "plot_2d",
-    "corner",
-    "plot_zz",
-    "plot_pp"
-    #    "diagonal_constraint",
-    #    "hist1d",
-    #    "lower_constraint",
-    #    "plot_empirical_z_score",
-    #    "empirical_z_score_corner",
-    #    "violin",
-]
+__all__ = ["plot_corner", "plot_zz", "plot_pp", "plot_posterior", "plot_pair"]

swyft/plot/constraint.py → swyft/plot/deprecated/constraint.py

@@ -1,12 +1,21 @@
 import matplotlib.patches as mpatches
 import matplotlib.path as mpath
 
-from swyft.plot.histogram import split_corner_axes
 from swyft.utils.marginals import get_d_dim_marginal_indices
 
+import numpy as np
+from typing import Tuple
+
+
+def _split_corner_axes(axes: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    diag = np.diag(axes)
+    lower = axes[np.tril(axes, -1).nonzero()]
+    upper = axes[np.triu(axes, 1).nonzero()]
+    return lower, diag, upper
+
 
 def diagonal_constraint(axes, bounds, alpha=0.25):
-    _, diag, _ = split_corner_axes(axes)
+    _, diag, _ = _split_corner_axes(axes)
     for i, ax in enumerate(diag):
         xlim = ax.get_xlim()
         ax.axvspan(xlim[0], bounds[i, 0], alpha=alpha)