quality_measures_128px_mask.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "146b5009-fad1-431f-b8b0-5a256aac54ca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import Pk_library as PKL\n",
    "import imageio\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import utilities\n",
    "import flow_architecture\n",
    "import losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fb50d8d-e848-4c44-85d8-c756b70965fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda'\n",
    "float_dtype = np.float32\n",
    "torch.set_default_tensor_type(torch.cuda.FloatTensor)\n",
    "device_id = 0\n",
    "torch.cuda.set_device(device_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55ba5420-4056-41cf-8062-b04bc36de001",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_dir = \"nbody_128px_mask_1p0_flow/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77d091ef-11a2-4298-94d8-d59eb3d84004",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Parameters():\n",
    "    def __init__(self):\n",
    "        #Data parameters\n",
    "        self.nx = 128\n",
    "        \n",
    "        #Fitting parameters\n",
    "        self.nlev_t = 1.0\n",
    "        self.noise_fac = self.nlev_t\n",
    "        self.noise_pix = 2*(self.nlev_t)**2\n",
    "        self.use_ql = False #The nbody power spectrum is matched with trainingdata\n",
    "        self.wf_batch_size = 100 #The number of maps to fit\n",
    "        mask128 = (imageio.imread(\"masks/mask1_128.png\")[:, :, 0]/255).astype(float)\n",
    "        self.mask = mask128 #np.ones((self.nx, self.nx))\n",
    "        \n",
    "        #Pre-trained flow parameters\n",
    "        self.flow_n_layers = 16\n",
    "        self.flow_hidden = [12, 12]\n",
    "        \n",
    "params = Parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e93102e2-1135-4f26-be2b-282539e7ebc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true_np = np.load(save_dir + 'true_maps.npy')\n",
    "y_pred_np = np.load(save_dir + 'masked_maps.npy')\n",
    "y_pred_flow_cc_np = np.load(save_dir + 'flow_maps.npy')\n",
    "y_pred_wf_cc_np = np.load(save_dir + 'wf_maps.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07a05b76-6a29-491c-848b-6f59f693783f",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = y_true_np.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3864ebd-c31c-47d8-a212-58314811ae62",
   "metadata": {},
   "outputs": [],
   "source": [
    "grid    = 128     #the map will have grid^2 pixels\n",
    "BoxSize = 512.0   #Mpc/h\n",
    "MAS     = 'None'  #MAS used to create the image; 'NGP', 'CIC', 'TSC', 'PCS' o 'None'\n",
    "threads = 1       #number of openmp threads"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc3c22d4-c591-4e60-aa53-8ad46f9c86d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "Pk_flow_mean  = 0\n",
    "Pk_flow_cc_mean = 0\n",
    "Pk_wf_mean = 0\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i in range(batch_size):\n",
    "        Pk2D_flow = PKL.Pk_plane(y_pred_flow_cc_np[i], BoxSize, MAS, threads, verbose=False)\n",
    "        Pk_flow_mean += Pk2D_flow.Pk\n",
    "        Pk2D_wf = PKL.Pk_plane(y_pred_wf_cc_np[i], BoxSize, MAS, threads, verbose=False)\n",
    "        Pk_wf_mean += Pk2D_wf.Pk\n",
    "\n",
    "    Pk_flow_mean /= params.wf_batch_size\n",
    "    Pk_wf_mean /= params.wf_batch_size\n",
    "    k = Pk2D_flow.k\n",
    "\n",
    "plt.plot(k[:-1], Pk_flow_mean[:-1], label='Flow')\n",
    "plt.plot(k[:-1], Pk_wf_mean[:-1], label='Wiener filtering')\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67fe2adc-958a-49ba-962b-de1e5bac7960",
   "metadata": {},
   "outputs": [],
   "source": [
    "fft_split_factor = 7.\n",
    "\n",
    "y_pred_flow_fft = np.fft.fft2(y_pred_flow_cc_np)\n",
    "y_pred_flow_fft_high = utilities.high_of_fft(y_pred_flow_fft, fft_split_factor)\n",
    "y_pred_wf_fft = np.fft.fft2(y_pred_wf_cc_np)\n",
    "y_pred_wf_fft_low = utilities.low_of_fft(y_pred_wf_fft, fft_split_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bae970ec-fd29-4ee6-8002-306bf9439775",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_split_fft = y_pred_flow_fft_high + y_pred_wf_fft_low\n",
    "y_pred_split = np.fft.ifft2(y_pred_split_fft).real"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "373f1e59-4c6d-4d0b-8caf-336b0f0585e2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Our simulations originally had the following std and mean\n",
    "sim_std = 14.3145\n",
    "sim_mean = 2.23550"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5717ce62-3d96-43a1-88ed-a29b82283e0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_overdensity_from_0mean_1std(array, sim_std, sim_mean):\n",
    "    \"\"\" Input array is 0 mean, 1 std. data.  It is returned to original simulation normalization, then normalized as an overdensity\"\"\"\n",
    "    array = (array + sim_mean)*sim_std\n",
    "    overdensity = (array - np.mean(array))/np.mean(array)\n",
    "    return overdensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54c5f1b9-ad66-4c98-beb4-537ede5e053f",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_true_np = make_overdensity_from_0mean_1std(y_true_np, sim_std, sim_mean)\n",
    "y_pred_np = make_overdensity_from_0mean_1std(y_pred_np, sim_std, sim_mean)\n",
    "y_pred_flow_cc_np = make_overdensity_from_0mean_1std(y_pred_flow_cc_np, sim_std, sim_mean)\n",
    "y_pred_wf_cc_np = make_overdensity_from_0mean_1std(y_pred_wf_cc_np, sim_std, sim_mean)\n",
    "y_pred_split = make_overdensity_from_0mean_1std(y_pred_split, sim_std, sim_mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c9e5d2-d306-42d1-a4e0-478f4e211aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "vmin = -0.5\n",
    "vmax = 4.5\n",
    "figsize = (6, 6)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5efc027e-9020-4c8f-97d4-9060614d3778",
   "metadata": {},
   "outputs": [],
   "source": [
    "utilities.imshow(y_true_np[0], vmin=vmin, vmax=vmax, figsize=figsize, axis=False, colorbar=False, file_name=save_dir+'truth.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87af32fc-53f2-4818-ad21-08508a35560c",
   "metadata": {},
   "outputs": [],
   "source": [
    "utilities.imshow(y_pred_np[0], vmin=vmin, vmax=vmax, figsize=figsize, axis=False, colorbar=False, file_name=save_dir+'masked.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0794c32-6782-4a71-a60b-6d304fb30077",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_np_vmin_mask = y_pred_np[0] - (1-params.mask)*100\n",
    "utilities.imshow(y_pred_np_vmin_mask, vmin=vmin, vmax=vmax, figsize=figsize, axis=False, colorbar=False, file_name=save_dir+'masked_2.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efdccb5b-b0d0-46c9-b10a-ae470e659e0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_np_vmin_mask = y_pred_np[0] * params.mask\n",
    "\n",
    "y_pred_np_vmin_mask = np.ma.masked_where(y_pred_np_vmin_mask == 0., y_pred_np_vmin_mask)\n",
    "\n",
    "plt.figure(figsize=figsize)\n",
    "cmap = matplotlib.cm.viridis\n",
    "cmap.set_bad('lightgreen', 1.)\n",
    "plt.imshow(y_pred_np_vmin_mask, vmin=vmin, vmax=vmax, cmap=cmap)\n",
    "plt.axis('off')\n",
    "plt.tight_layout()\n",
    "plt.savefig(save_dir+'masked_3.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "547547da-7e70-47c6-9433-ee6fff1c08a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "utilities.imshow(y_pred_wf_cc_np[0], vmin=vmin, vmax=vmax, figsize=figsize, axis=False, colorbar=False, file_name=save_dir+'wf.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bcba5c9-006a-4dea-a7b4-c87ad914e164",
   "metadata": {},
   "source": [
    "Fourier split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16eb21f3-8a27-44a4-828a-c38bb7e6f0a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "utilities.imshow(y_pred_split[0], vmin=vmin, vmax=vmax, colorbar=False, figsize=figsize, axis=False, file_name=save_dir+'flow.png')\n",
    "\n",
    "\"\"\"J2_split_mean = 0.\n",
    "for n in range(params.wf_batch_size):\n",
    "    J2_split_mean += lossfunctions.loss_J2(y_true[n], torch.tensor(y_pred_split[n])).cpu().detach().numpy()\n",
    "J2_split_mean = J2_split_mean / params.wf_batch_size\n",
    "print(J2_split_mean)\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5753dc5-2e73-47f6-a7f8-52780c750a27",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.mean((y_pred_split.astype('float32') - y_true_np.astype('float32'))**2))\n",
    "print(np.mean((y_pred_wf_cc_np.astype('float32') - y_true_np.astype('float32'))**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5764217-7c1a-42b4-b71f-cd71bc51952c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.mean(((y_pred_split.astype('float32') - y_true_np.astype('float32'))*params.mask)**2))\n",
    "print(np.mean(((y_pred_wf_cc_np.astype('float32') - y_true_np.astype('float32'))*params.mask)**2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fc2a881-3890-480f-b0df-496331ddcd82",
   "metadata": {},
   "outputs": [],
   "source": [
    "interior_frac = np.sum(params.mask) / (params.nx**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d6f177-e12a-4ab0-95d1-986e45ede70c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.mean(((y_pred_split.astype('float32') - y_true_np.astype('float32'))*params.mask)**2) / interior_frac)\n",
    "print(np.mean(((y_pred_wf_cc_np.astype('float32') - y_true_np.astype('float32'))*params.mask)**2) / interior_frac)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0d2440d-0b00-4411-bcd7-cc44e0d3f76d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 14})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a84a6ab8-5659-409a-aaaf-4904e806077f",
   "metadata": {},
   "outputs": [],
   "source": [
    "Pk_true = np.zeros((batch_size, 90))\n",
    "Pk_flow = np.zeros((batch_size, 90))\n",
    "Pk_flow_cc = np.zeros((batch_size, 90))\n",
    "Pk_wf = np.zeros((batch_size, 90))\n",
    "Pk_wf_cc = np.zeros((batch_size, 90))\n",
    "Pk_noise = np.zeros((batch_size, 90))\n",
    "N_flow = np.zeros((batch_size, 90))\n",
    "N_wf = np.zeros((batch_size, 90))\n",
    "\n",
    "with torch.no_grad():\n",
    "    for i in range(batch_size):\n",
    "        Pk2D_flow_true = PKL.XPk_plane((y_pred_split[i]*params.mask).astype('float32'), (y_true_np[i]*params.mask).astype('float32'), BoxSize, MAS, MAS, threads)\n",
    "        Pk_flow[i] = Pk2D_flow_true.Pk[:, 0]\n",
    "        Pk_true[i] = Pk2D_flow_true.Pk[:, 1]\n",
    "        Pk_flow_cc[i] = Pk2D_flow_true.r\n",
    "        \n",
    "        Pk2D_wf_true = PKL.XPk_plane((y_pred_wf_cc_np[i]*params.mask).astype('float32'), (y_true_np[i]*params.mask).astype('float32'), BoxSize, MAS, MAS, threads)\n",
    "        Pk_wf[i] = Pk2D_wf_true.Pk[:, 0]\n",
    "        Pk_wf_cc[i] = Pk2D_wf_true.r\n",
    "        \n",
    "        noise = y_pred_np[i] - y_true_np[i]\n",
    "        Pk2D_noise = PKL.Pk_plane((noise*params.mask).astype('float32'), BoxSize, MAS, threads)\n",
    "        Pk_noise[i] = Pk2D_noise.Pk\n",
    "        \n",
    "        epsilon_flow = y_pred_split[i] - y_true_np[i]\n",
    "        Pk2D_Nflow = PKL.Pk_plane((epsilon_flow*params.mask).astype('float32'), BoxSize, MAS, threads)\n",
    "        N_flow[i] = Pk2D_Nflow.Pk\n",
    "        \n",
    "        epsilon_wf = y_pred_wf_cc_np[i] - y_true_np[i]\n",
    "        Pk2D_Nwf = PKL.Pk_plane((epsilon_wf*params.mask).astype('float32'), BoxSize, MAS, threads)\n",
    "        N_wf[i] = Pk2D_Nwf.Pk\n",
    "        \n",
    "    kvals = Pk2D_flow_true.k\n",
    "\n",
    "\n",
    "rcut = 4\n",
    "\n",
    "kvals = kvals[:-rcut]\n",
    "    \n",
    "Pk_flow_cc_mean = np.mean(Pk_flow_cc[:, :-rcut], 0)\n",
    "Pk_flow_cc_1sigma = np.std(Pk_flow_cc[:, :-rcut], 0)\n",
    "Pk_wf_cc_mean = np.mean(Pk_wf_cc[:, :-rcut], 0)\n",
    "Pk_wf_cc_1sigma = np.std(Pk_wf_cc[:, :-rcut], 0)\n",
    "\n",
    "Pk_figsize = (5, 4)\n",
    "plt.figure(figsize=Pk_figsize)\n",
    "plt.plot(kvals, Pk_flow_cc_mean, label=r'$r^{\\mathrm{true, flow}}$')\n",
    "plt.fill_between(kvals, Pk_flow_cc_mean+Pk_flow_cc_1sigma, Pk_flow_cc_mean-Pk_flow_cc_1sigma, alpha=0.2)\n",
    "plt.plot(kvals, Pk_wf_cc_mean, label=r'$r^{\\mathrm{true, Wiener\\ filtered}}$')\n",
    "plt.fill_between(kvals, Pk_wf_cc_mean+Pk_wf_cc_1sigma, Pk_wf_cc_mean-Pk_wf_cc_1sigma, alpha=0.2)\n",
    "plt.xscale('log')\n",
    "plt.xlabel(r'$k\\ (h/\\mathrm{Mpc})$')\n",
    "plt.ylabel(r'$r(k)$')\n",
    "plt.legend(loc='lower left')\n",
    "#plt.tight_layout()\n",
    "plt.savefig(save_dir + 'r_0p1noise_mask_small.pdf', bbox_inches='tight')\n",
    "plt.show()\n",
    "\n",
    "\n",
    "Pk_true_mean = np.mean(Pk_true[:, :-rcut], 0)\n",
    "Pk_true_1sigma = np.std(Pk_true[:, :-rcut], 0)\n",
    "Pk_flow_mean = np.mean(Pk_flow[:, :-rcut], 0)\n",
    "Pk_flow_1sigma = np.std(Pk_flow[:, :-rcut], 0)\n",
    "Pk_wf_mean = np.mean(Pk_wf[:, :-rcut], 0)\n",
    "Pk_wf_1sigma = np.std(Pk_wf[:, :-rcut], 0)\n",
    "Pk_noise_mean = np.mean(Pk_noise[:, :-rcut], 0)\n",
    "Pk_noise_1sigma = np.std(Pk_noise[:, :-rcut], 0)\n",
    "N_flow_mean = np.mean(N_flow[:, :-rcut], 0)\n",
    "N_flow_1sigma = np.std(N_flow[:, :-rcut], 0)\n",
    "N_wf_mean = np.mean(N_wf[:, :-rcut], 0)\n",
    "N_wf_1sigma = np.std(N_wf[:, :-rcut], 0)\n",
    "\n",
    "plt.figure(figsize=Pk_figsize)\n",
    "plt.plot(kvals, Pk_flow_mean, label='Flow posterior')\n",
    "plt.fill_between(kvals, Pk_flow_mean+Pk_flow_1sigma, Pk_flow_mean-Pk_flow_1sigma, alpha=0.2)\n",
    "plt.plot(kvals, Pk_wf_mean, label='Wiener filtered')\n",
    "plt.fill_between(kvals, Pk_wf_mean+Pk_wf_1sigma, Pk_wf_mean-Pk_wf_1sigma, alpha=0.2)\n",
    "plt.plot(kvals, Pk_true_mean, label='Truth', linestyle='--')\n",
    "plt.fill_between(kvals, Pk_true_mean+Pk_true_1sigma, Pk_true_mean-Pk_true_1sigma, alpha=0.2)\n",
    "plt.plot(kvals, Pk_noise_mean, label=r'$1.0\\tilde{\\sigma}\\ \\mathrm{noise}$')\n",
    "plt.fill_between(kvals, Pk_noise_mean+Pk_noise_1sigma, Pk_noise_mean-Pk_noise_1sigma, alpha=0.2)\n",
    "plt.plot(kvals, N_flow_mean, label=r'$N^{\\mathrm{flow}}$')\n",
    "plt.fill_between(kvals, N_flow_mean+N_flow_1sigma, N_flow_mean-N_flow_1sigma, alpha=0.2)\n",
    "plt.plot(kvals, N_wf_mean, label=r'$N^{\\mathrm{Wiener\\ filtered}}$')\n",
    "plt.fill_between(kvals, N_wf_mean+N_wf_1sigma, N_wf_mean-N_wf_1sigma, alpha=0.2)\n",
    "plt.xscale('log')\n",
    "plt.yscale('log')\n",
    "plt.xlabel(r'$k\\ (h/\\mathrm{Mpc})$')\n",
    "plt.ylabel(r'$P(k)\\ (\\mathrm{Mpc}^2/h^2)$')\n",
    "plt.legend(loc='lower left')\n",
    "#plt.tight_layout()\n",
    "plt.savefig(save_dir + 'ps_0p1noise_mask_small.pdf', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46fdc08f-2cd9-4433-aa44-18be3e58bbc9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}