Updated config space paper plots

cosmo_inference/notebooks/2D_cosmic_shear_configuration_plots/get_prior_psf_leakage.ipynb

@@ -0,0 +1,249 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "2978bafe",
+   "metadata": {},
+   "source": [
+    "# Covariance matrix and PSF leakage\n",
+    "\n",
+    "This notebook plots the combined covariance matrix, and samples and plots the 2D marginalised posteriors of the PSF leakage parameters $\\alpha$ and $\\beta$."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "54601eff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "if not os.path.exists(\"./Plots\"):\n",
+    "    os.makedirs(\"./Plots\")\n",
+    "\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "from astropy.io import fits\n",
+    "import seaborn as sns\n",
+    "from getdist import plots, loadMCSamples, MCSamples\n",
+    "\n",
+    "from shear_psf_leakage.rho_tau_stat import RhoStat, TauStat, PSFErrorFit\n",
+    "\n",
+    "# Use paper style and seaborn with husl palette\n",
+    "plt.style.use(\n",
+    "    \"/home/guerrini/matplotlib_config/paper.mplstyle\"\n",
+    ")\n",
+    "# Set default palette - will be updated per plot as needed\n",
+    "sns.set_palette(\"husl\")\n",
+    "%matplotlib inline\n",
+    "\n",
+    "g = plots.get_subplot_plotter(width_inch=30)\n",
+    "g.settings.axes_fontsize=30\n",
+    "g.settings.axes_labelsize=30\n",
+    "g.settings.alpha_filled_add = 0.7\n",
+    "g.settings.legend_fontsize = 25\n",
+    "\n",
+    "ver = 'v1.4.6.3'\n",
+    "blind = 'B'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4a1618a7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_path = f\"/home/guerrini/sp_validation/cosmo_inference/data/SP_{ver}_config/\"\n",
+    "\n",
+    "path_cosmo_val = \"/home/guerrini/sp_validation/notebooks/cosmo_val/output/\"\n",
+    "\n",
+    "roots = [\n",
+    "    f\"SP_{ver}_{blind}\",\n",
+    "    f\"SP_{ver}_leak_corr_{blind}\"\n",
+    "]\n",
+    "\n",
+    "labels = [\n",
+    "    f\"SP_{ver}_{blind}\",\n",
+    "    f\"SP_{ver}_leak_corr_{blind}\"\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6c8c0a1f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_vectors = []\n",
+    "\n",
+    "for root in roots:\n",
+    "    data_vectors.append(\n",
+    "        fits.open(data_path  +f\"SP_{ver}_{blind}/cosmosis_{root}_masked.fits\")\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "af22339d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def cov_to_corr(cov):\n",
+    "    \"\"\"Convert a covariance matrix to a correlation matrix.\"\"\"\n",
+    "    d = np.sqrt(np.diag(cov))\n",
+    "    corr = cov / np.outer(d, d)\n",
+    "    corr[cov == 0] = 0\n",
+    "    return corr"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a5d8de0d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Print the covariance matrix for each root\n",
+    "for i, root in enumerate(roots):\n",
+    "    print(f\"Covariance matrix for {labels[i]}:\")\n",
+    "    cov = data_vectors[i][\"COVMAT\"].data\n",
+    "\n",
+    "    n_bins = cov.shape[0] // 4\n",
+    "\n",
+    "    fig, ax = plt.subplots(figsize=(10, 8))\n",
+    "\n",
+    "    im = ax.imshow(cov_to_corr(cov), vmin=-1, vmax=1, cmap=\"seismic\")\n",
+    "    ax.set_aspect('equal')\n",
+    "    ax.set_yticks(np.array([10, 30, 50, 70]))\n",
+    "    ax.set_yticklabels([r\"$\\xi_+(\\vartheta)$\", r\"$\\xi_-(\\vartheta)$\", r\"$\\tau_0(\\vartheta)$\", r\"$\\tau_2(\\vartheta)$\"])\n",
+    "    ax.set_xticks(np.array([10, 30, 50, 70]))\n",
+    "    ax.set_xticklabels([r\"$\\xi_+(\\vartheta)$\", r\"$\\xi_-(\\vartheta)$\", r\"$\\tau_0(\\vartheta)$\", r\"$\\tau_2(\\vartheta)$\"], rotation=45)\n",
+    "    fig.colorbar(im, ax=ax)\n",
+    "\n",
+    "    plt.savefig(f\"./Plots/cov_matrix_{root}.png\", bbox_inches='tight', dpi=300)\n",
+    "    plt.show()\n",
+    "    print(\"\\n\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a484afb6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Create dummy rho and tau stat handler.\n",
+    "\n",
+    "#Inference of the xi_sys parameters\n",
+    "sep_units = 'arcmin'\n",
+    "coord_units = 'degrees'\n",
+    "theta_min = 1.0\n",
+    "theta_max = 250\n",
+    "nbins = 20\n",
+    "\n",
+    "\n",
+    "TreeCorrConfig_xi = {\n",
+    "    'ra_units': coord_units,\n",
+    "    'dec_units': coord_units,\n",
+    "    'min_sep': theta_min,\n",
+    "    'max_sep': theta_max,\n",
+    "    'sep_units': sep_units,\n",
+    "    'nbins': nbins,\n",
+    "    'var_method':'jackknife',\n",
+    "}\n",
+    "\n",
+    "rho_stats_handler = RhoStat(\n",
+    "    output='.',\n",
+    "    treecorr_config=TreeCorrConfig_xi,\n",
+    "    verbose=True\n",
+    ")\n",
+    "\n",
+    "tau_stats_handler = TauStat(\n",
+    "    catalogs=rho_stats_handler.catalogs,\n",
+    "    output='.',\n",
+    "    treecorr_config=TreeCorrConfig_xi,\n",
+    "    verbose=True\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a4dd4416",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#Create a PSFErrorFit instance\n",
+    "psf_fitter = PSFErrorFit(rho_stats_handler, tau_stats_handler, path_cosmo_val+'rho_tau_stats/', use_eta=False)\n",
+    "\n",
+    "g  = plots.get_subplot_plotter(width_inch=30)\n",
+    "\n",
+    "g.settings.axes_fontsize=30\n",
+    "g.settings.axes_labelsize=30\n",
+    "g.settings.alpha_filled_add = 0.7\n",
+    "g.settings.legend_fontsize = 40\n",
+    "\n",
+    "chains = []\n",
+    "\n",
+    "#Load rho-, tau-statistics, and cov_tau from the data_vector\n",
+    "for i, root in enumerate(roots):\n",
+    "    print(\"Sampling PSF parameters for \", labels[i])\n",
+    "    path_rho = f\"rho_stats_{root}.fits\"\n",
+    "    path_tau = f\"tau_stats_{root}.fits\"\n",
+    "    path_cov_rho = f\"cov_rho_{root}.npy\"\n",
+    "    path_cov_tau = f\"cov_tau_{root}_th.npy\"\n",
+    "    psf_fitter.load_rho_stat(path_rho)\n",
+    "    psf_fitter.load_tau_stat(path_tau)\n",
+    "    psf_fitter.load_covariance(path_cov_rho, cov_type='rho')\n",
+    "    psf_fitter.load_covariance(path_cov_tau, cov_type='tau')\n",
+    "    samples_lq, _, _ = psf_fitter.get_least_squares_params_samples(npatch=None, apply_debias=False)\n",
+    "\n",
+    "    samples_gd = MCSamples(samples=samples_lq, names=[r\"\\alpha\", r\"\\beta\"], labels=[r\"\\alpha\", r\"\\beta\"])\n",
+    "\n",
+    "    chains.append(samples_gd)\n",
+    "\n",
+    "g.triangle_plot(chains,\n",
+    "                filled=True,\n",
+    "                legend_labels=labels,\n",
+    "                legend_loc='upper right',\n",
+    "                \n",
+    ")\n",
+    "\n",
+    "# plt.savefig(f\"./Plots/psf_leakage_params.png\", bbox_inches='tight', dpi=300)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a24f3510",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "my_env",
+   "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.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}