Add plots for evaluation that only use likelihood (#54)

* move notebooks around * add optax * notebook * add posterior on |g| * add outlier figures * draf tnotebooks with some checks * small correction
LSSTDESC · Dec 4, 2024 · cf26964 · cf26964
1 parent 2ada016
commit cf26964
Show file tree

Hide file tree

Showing 42 changed files with 1,395 additions and 15 deletions.
diff --git a/experiments/exp1/figs/calibration.pdf b/experiments/exp1/figs/calibration.pdf
diff --git a/experiments/exp1/figs/contours.pdf b/experiments/exp1/figs/contours.pdf
diff --git a/experiments/exp1/figs/mean_std_hist.pdf b/experiments/exp1/figs/mean_std_hist.pdf
diff --git a/experiments/exp1/figs/multiplicative_bias_hist.pdf b/experiments/exp1/figs/multiplicative_bias_hist.pdf
diff --git a/experiments/exp1/figs/traces.pdf b/experiments/exp1/figs/traces.pdf
diff --git a/experiments/exp1/make_figures.py b/experiments/exp1/make_figures.py
@@ -6,7 +6,6 @@
 os.environ["JAX_PLATFORMS"] = "cpu"
 os.environ["JAX_ENABLE_X64"] = "True"
 
-from math import sqrt
 
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
@@ -36,6 +35,9 @@ def make_trace_plots(g_samples: Array, n_examples: int = 25) -> None:
 
             ax1.plot(g1)
             ax2.plot(g2)
+            ax1.set_ylabel("g1")
+            ax2.set_ylabel("g2")
+
             ax1.set_title(f"Index: {ii}")
 
             pdf.savefig(fig)
@@ -87,6 +89,7 @@ def make_histogram_mbias(g_samples: Array) -> None:
         mbias = (g1.mean(axis=1) - 0.02) / 0.02
         fig, ax = plt.subplots(1, 1, figsize=(7, 7))
         ax.hist(mbias, bins=31, histtype="step")
+        ax.set_xlabel("m")
         pdf.savefig(fig)
         plt.close(fig)
 
@@ -95,22 +98,65 @@ def make_histogram_means_and_stds(g_samples: Array) -> None:
     fname = "figs/mean_std_hist.pdf"
     with PdfPages(fname) as pdf:
         g1 = g_samples[:, :, 0]
+        g2 = g_samples[:, :, 1]
 
         means = g1.mean(axis=1)
         stds = g1.std(axis=1)
 
         fig, ax = plt.subplots(1, 1, figsize=(7, 7))
         ax.hist(means, bins=31, histtype="step")
         ax.set_title(f"Std: {means.std():.5g}")
+        ax.set_xlabel("Mean of posterior of g1")
+        ax.axvline(means.mean(), linestyle="--", color="k", label="mean")
+        ax.legend()
+        pdf.savefig(fig)
+        plt.close(fig)
+
+        fig, ax = plt.subplots(1, 1, figsize=(7, 7))
+        ax.hist(stds, bins=31, histtype="step")
+        ax.set_xlabel("Std of posterior of g1")
+        ax.axvline(stds.mean(), linestyle="--", color="k", label="mean")
+        ax.legend()
+        pdf.savefig(fig)
+        plt.close(fig)
+
+        means = g2.mean(axis=1)
+        stds = g2.std(axis=1)
+
+        fig, ax = plt.subplots(1, 1, figsize=(7, 7))
+        ax.hist(means, bins=31, histtype="step")
+        ax.set_title(f"Std: {means.std():.5g}")
+        ax.set_xlabel("Mean of posterior of g2")
         ax.axvline(means.mean(), linestyle="--", color="k", label="mean")
         ax.legend()
         pdf.savefig(fig)
         plt.close(fig)
 
         fig, ax = plt.subplots(1, 1, figsize=(7, 7))
         ax.hist(stds, bins=31, histtype="step")
-        ax.axvline(1e-3 / sqrt(1000), linestyle="--", color="k")
-        ax.set_title(f"Std_correct: {1e-3 / sqrt(1000) / sqrt(2):.5g}")
+        ax.set_xlabel("Std of posterior of g2")
+        ax.axvline(stds.mean(), linestyle="--", color="k", label="mean")
+        ax.legend()
+        pdf.savefig(fig)
+        plt.close(fig)
+
+        g_mag = jnp.sqrt(g1**2 + g2**2)
+        means = g_mag.mean(axis=1)
+        stds = g_mag.std(axis=1)
+        fig, ax = plt.subplots(1, 1, figsize=(7, 7))
+        ax.hist(means, bins=31, histtype="step")
+        ax.set_xlabel("Mean of posterior of |g|")
+        ax.set_title(f"Std: {means.std():.5g}")
+        ax.axvline(means.mean(), linestyle="--", color="k", label="mean")
+        ax.legend()
+        pdf.savefig(fig)
+        plt.close(fig)
+
+        fig, ax = plt.subplots(1, 1, figsize=(7, 7))
+        ax.hist(stds, bins=31, histtype="step")
+        ax.set_xlabel("Std of posterior of |g|")
+        ax.axvline(stds.mean(), linestyle="--", color="k", label="mean")
+        ax.legend()
         pdf.savefig(fig)
         plt.close(fig)
 

diff --git a/experiments/exp2/figs/contours.pdf b/experiments/exp2/figs/contours.pdf
diff --git a/experiments/exp2/figs/convergence_hist.pdf b/experiments/exp2/figs/convergence_hist.pdf
diff --git a/experiments/exp2/figs/timing.pdf b/experiments/exp2/figs/timing.pdf
diff --git a/experiments/exp2/figs/traces.pdf b/experiments/exp2/figs/traces.pdf
diff --git a/experiments/exp2/figs/traces_adapt.pdf b/experiments/exp2/figs/traces_adapt.pdf
diff --git a/experiments/exp2/figs/traces_out.pdf b/experiments/exp2/figs/traces_out.pdf
diff --git a/experiments/exp2/make_figures.py b/experiments/exp2/make_figures.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import os
+from typing import Iterable
 
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 os.environ["JAX_PLATFORMS"] = "cpu"
@@ -25,15 +26,22 @@ def make_trace_plots(
     samples_dict: dict[str, Array],
     truth: dict[str, Array],
     fpath: str,
-    n_examples: int = 25,
+    n_examples: int = 50,
 ) -> None:
     """Make example figure showing example trace plots for each parameter."""
-    # by default, we choose 10 random traces to plot in 1 PDF file.
     n_gals, _, _ = samples_dict["lf"].shape
+    indices = np.random.choice(np.arange(0, n_gals), size=(n_examples,), replace=False)
+    make_trace_at_indices(indices, samples_dict, truth, fpath)
 
+
+def make_trace_at_indices(
+    indices: Iterable,
+    samples_dict: dict[str, Array],
+    truth: dict[str, Array],
+    fpath: str,
+):
     with PdfPages(fpath) as pdf:
-        for _ in tqdm(range(n_examples), desc="Making traces"):
-            idx = np.random.choice(np.arange(0, n_gals)).item()
+        for idx in tqdm(indices, desc="Making traces"):
             chains = {k: v[idx] for k, v in samples_dict.items()}
             tv = {p: q[idx].item() for p, q in truth.items()}
 
@@ -93,12 +101,16 @@ def make_convergence_histograms(samples_dict: dict[str, Array]) -> None:
             ess[p].append(effective_sample_size(chains) / n_samples_total)
 
     # print rhat outliers
+    outliers_indices = set()
     for p in rhats:
         rhat = np.array(rhats[p])
         _ess = np.array(ess[p])
-        n_outliers = sum((rhat < 0.99) | (rhat > 1.05) | (_ess < 0.25))
+        outliers = (rhat < 0.99) | (rhat > 1.05) | (_ess < 0.25)
+        n_outliers = sum(outliers)
         with open("figs/outliers.txt", "a", encoding="utf-8") as f:
             print(f"Number of R-hat outliers for {p}: {n_outliers}", file=f)
+        indices = np.argwhere(outliers)
+        outliers_indices = outliers_indices.union(set(indices.ravel()))
 
     with PdfPages(fname) as pdf:
         for p in samples_dict:
@@ -117,6 +129,8 @@ def make_convergence_histograms(samples_dict: dict[str, Array]) -> None:
             pdf.savefig(fig)
             plt.close(fig)
 
+    return outliers_indices
+
 
 def make_timing_plots(results_dict: dict) -> None:
     print("INFO: Making timing plots...")
@@ -177,12 +191,19 @@ def main():
     # make plots
     make_trace_plots(samples, truth, fpath="figs/traces.pdf")
     make_contour_plots(samples, truth)
-    make_convergence_histograms(samples)
+    out_indices = make_convergence_histograms(samples)
     make_timing_plots(results)
 
     # on adaption too
     adapt_states = results[max_n_gal]["adapt_info"].state.position
-    make_trace_plots(adapt_states, truth, fpath="figs/traces_adapt.pdf", n_examples=50)
+    make_trace_plots(adapt_states, truth, fpath="figs/traces_adapt.pdf")
+
+    # outliers
+    if len(out_indices) > 0:
+        make_trace_at_indices(out_indices, samples, truth, fpath="figs/traces_out.pdf")
+    else:  # avoid confusion with previous
+        if Path("figs/traces_out.pdf").exists():
+            os.remove("figs/traces_out.pdf")
 
 
 if __name__ == "__main__":

diff --git a/notebooks/check-g-mag-random.ipynb b/notebooks/check-g-mag-random.ipynb
@@ -0,0 +1,170 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "rng = np.random.RandomState(43)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "g1, g2 = rng.normal(size=(2, 1_000_000))*1e-3 + 0.02"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([[1.00132769e-06, 6.64097921e-10],\n",
+       "       [6.64097921e-10, 9.97662766e-07]])"
+      ]
+     },
+     "execution_count": 49,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.cov(g1, g2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(0.0010006631265643401, 0.000998830200162854, 0.0009997776905801786)"
+      ]
+     },
+     "execution_count": 50,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.std(g1), np.std(g2), np.std(np.sqrt(g1**2 + g2**2)) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(1.6044601312014296e-09,\n",
+       " 1.5979783261011428e-09,\n",
+       " 3.204493672957289e-09,\n",
+       " 9.995554305818354e-07)"
+      ]
+     },
+     "execution_count": 51,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.var(g1**2), np.var(g2**2), np.var(g1**2 + g2**2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# g1+= 0.02"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(1.4332765422958652e-06, 1.4438625190160518e-06, 2.032142858328091e-06)"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "4.019049885896081e-05"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "np.std(g1**2 + g2**2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "bpd_gpu2",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}