{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook is part of the SelfiSys project and is licensed under the GNU General Public License v3.0 or later (GPL-3.0-or-later). See the accompanying LICENSE file or visit [GNU GPL](https://www.gnu.org/licenses/gpl-3.0.en.html) for details.\n",
    "\n",
    "Tristan Hoellinger<br/>\n",
    "Institut d'Astrophysique de Paris</br>\n",
    "tristan.hoellinger@iap.fr\n",
    "\n",
    "# Survey mask\n",
    "\n",
    "Generate a survey mask from an external sky map."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the environment and parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from selfisys.global_parameters import *\n",
    "\n",
    "plotly = False\n",
    "wd = Path(OUTPUT_PATH) / \"expl_notebooks/surveymask/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import exists\n",
    "from datetime import datetime\n",
    "import warnings\n",
    "\n",
    "import healpy as hp\n",
    "from healpy.newvisufunc import projview, newprojplot\n",
    "from matplotlib import colormaps\n",
    "from matplotlib.colors import LinearSegmentedColormap, rgb2hex\n",
    "\n",
    "# Project-specific imports\n",
    "from selfisys.utils.plot_params import *  # Plotting configuration\n",
    "from selfisys.utils.plot_utils import *  # Utility functions for plotting\n",
    "if plotly:\n",
    "    from selfisys.utils.examples_utils import clear_large_plot\n",
    "\n",
    "# Use today's date string for labelling outputs\n",
    "today = datetime.today().strftime(\"%b%d\")\n",
    "output_dir = Path(wd) / \"expl\"\n",
    "output_dir.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the base map\n",
    "hpx = hp.read_map(Path(OUTPUT_PATH) / \"DATA/external_map.fits\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate the 3D survey mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 64  # Number of grid points\n",
    "L = 3.6  # Box size (in Gpc/h)\n",
    "force = True  # Force re-computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rotation angles in Galactic coordinates for map adjustment\n",
    "rotation = np.array([80, 20])  # [latitude, longitude]\n",
    "\n",
    "# Define origin in our custom convention (relative to Healpy's default)\n",
    "c = np.array([np.pi / 2, 0])  # [latitude, longitude]\n",
    "\n",
    "# Custom colormap\n",
    "cmap_blue = colormaps.get_cmap(\"Blues\")\n",
    "custom_cm = LinearSegmentedColormap.from_list(\n",
    "    \"custom_cm\", [(1, 1, 1), rgb2hex(cmap_blue(cmap_blue.N))], N=100\n",
    ")\n",
    "\n",
    "# Rotate the origin\n",
    "lat_origin_galactic, lon_origin_galactic = c\n",
    "rotator = hp.Rotator(rot=rotation)\n",
    "rot_lat_origin_galactic, rot_lon_origin_galactic = rotator(\n",
    "    lat_origin_galactic, lon_origin_galactic\n",
    ")\n",
    "\n",
    "# Define lines for visualisation\n",
    "ndots = 30  # Number of points per line\n",
    "line1_theta = c[0] + np.linspace(-np.pi / 2, 0, ndots)\n",
    "line1_phi = c[1] + [np.pi / 4] * ndots\n",
    "line2_theta = c[0] + [0] * ndots\n",
    "line2_phi = c[1] + np.linspace(-np.pi / 4, np.pi / 4, ndots)\n",
    "line3_theta = c[0] + np.linspace(-np.pi / 2, 0, ndots)\n",
    "line3_phi = c[1] + [-np.pi / 4] * ndots"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Raw mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Grid of Cartesian coordinates ranging from 0 to L\n",
    "x = np.linspace(0, L, N)\n",
    "xx, yy, zz = np.meshgrid(x, x, x)\n",
    "\n",
    "# Grid of radial distances\n",
    "grid_cart = np.sqrt(xx**2 + yy**2 + zz**2)\n",
    "\n",
    "# Check that the maximum radial distance matches the expected value, sqrt(3)*L\n",
    "assert np.isclose(grid_cart[-1, -1, -1], np.sqrt(3) * L)\n",
    "\n",
    "# Convert Cartesian coordinates to spherical coordinates (phi, theta, rho)\n",
    "grid_sph = np.zeros((N, N, N, 3))\n",
    "grid_sph[..., 0] = np.arccos(zz / (grid_cart + 1e-10))  # phi (colatitude)\n",
    "grid_sph[..., 1] = np.arctan2(yy, xx)  # theta (azimuthal angle)\n",
    "grid_sph[..., 2] = grid_cart  # rho (radial distance)\n",
    "\n",
    "# Convert to latitude, longitude, and radius\n",
    "grid_sph_latlon = np.zeros_like(grid_sph)\n",
    "grid_sph_latlon[..., 0] = np.pi / 2 - grid_sph[..., 0]  # latitude\n",
    "grid_sph_latlon[..., 1] = grid_sph[..., 1] - np.pi / 4  # longitude (shifted by 45°)\n",
    "grid_sph_latlon[..., 2] = grid_sph[..., 2]  # radius\n",
    "\n",
    "# Convert to degrees\n",
    "grid_sph_latlon_degrees = np.zeros_like(grid_sph)\n",
    "grid_sph_latlon_degrees[..., 0] = np.degrees(grid_sph_latlon[..., 0])  # latitude (degrees)\n",
    "grid_sph_latlon_degrees[..., 1] = np.degrees(grid_sph_latlon[..., 1])  # longitude (degrees)\n",
    "grid_sph_latlon_degrees[..., 2] = grid_sph_latlon[..., 2]  # radius (unchanged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = f\"raw_mask_N{N}.npy\"  # filename for the raw mask\n",
    "\n",
    "if not exists(wd / filename) or force:\n",
    "    # Rotate the HEALPix map from Galactic to custom coordinates\n",
    "    rotator = hp.rotator.Rotator(rot=rotation, coord=[\"C\", \"G\"], deg=True)\n",
    "    hpx_rotated = rotator.rotate_map_pixel(hpx)\n",
    "\n",
    "    # Compute latitudes and longitudes and interpolate\n",
    "    lats = c[0] - grid_sph_latlon[..., 0]\n",
    "    longs = c[1] + grid_sph_latlon[..., 1]\n",
    "    grid_mask = np.zeros((N, N, N))\n",
    "    grid_mask = hp.get_interp_val(hpx_rotated, lats, longs, lonlat=False)\n",
    "\n",
    "    np.save(wd / filename, grid_mask)\n",
    "else:\n",
    "    grid_mask = np.load(wd / filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualise the 3D grid mask using Plotly if enabled\n",
    "if plotly:\n",
    "    if N <= 128:\n",
    "        fig = plotly_3d(grid_mask, size=N, L=L, colormap=\"Blues\")\n",
    "    else:\n",
    "        # Downsample the grid for visualisation at higher resolutions\n",
    "        downsample_factor = N // 128\n",
    "        downsampled_field = grid_mask[\n",
    "            ::downsample_factor, ::downsample_factor, ::downsample_factor\n",
    "        ]\n",
    "        fig = plotly_3d(downsampled_field, L=L, colormap=\"Blues\")\n",
    "\n",
    "    fig.show()\n",
    "    # clear_large_plot(fig)  # Clear the Plotly figure to avoid memory issues"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mask with additional holes + linear extinction near the galactic plane"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "maxlat = 60 * np.pi / 180  # max latitude for linear extinction\n",
    "Nholes = 256  # number of additional holes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def inverse_rotation(lat, long):\n",
    "    \"\"\"\n",
    "    Inverse rotation of the galactic plane.\n",
    "    Take latitude and longitude in the rotated coordinate and return\n",
    "    the corresponding latitude and longitude in the galactic coordinate.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    lat : float or array-like\n",
    "        Latitude in the rotated coordinate system (radians).\n",
    "    long : float or array-like\n",
    "        Longitude in the rotated coordinate system (radians).\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    tuple\n",
    "        Corresponding latitude and longitude in galactic coordinate.\n",
    "    \"\"\"\n",
    "    return hp.Rotator(rot=rotation, inv=True)(lat, long)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(1)\n",
    "lats_centers_small_holes = np.random.uniform(0, np.pi / 2, 2048)\n",
    "longs_centers_small_holes = np.random.uniform(-np.pi / 4, np.pi / 4, 2048)\n",
    "lats_centers_small_holes = lats_centers_small_holes[:Nholes]\n",
    "longs_centers_small_holes = longs_centers_small_holes[:Nholes]\n",
    "x_idx = -1\n",
    "y_idx = -1\n",
    "z_idx = -1\n",
    "lat = grid_sph_latlon[x_idx, y_idx, z_idx, 0]\n",
    "long = grid_sph_latlon[x_idx, y_idx, z_idx, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 1301760 characters]",
      "text/plain": [
       "<Figure size 3000x1890 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "# Visualise the mask using a Mollweide projection\n",
    "projview(\n",
    "    hpx,\n",
    "    coord=[\"C\", \"G\"],\n",
    "    projection_type=\"mollweide\",\n",
    "    rot=rotation,\n",
    "    cmap=custom_cm,\n",
    "    override_rot_graticule_properties={\"g_linewidth\": 0.2, \"g_linestyle\": \"--\", \"g_color\": \"red\"},\n",
    "    graticule=False,\n",
    "    xlabel=r\"longitude $\\ell$ (°)\",\n",
    "    ylabel=r\"latitude $b$ (°)\",\n",
    "    cbar=True,\n",
    "    cbar_ticks=[0, 0.5, 1],\n",
    "    cb_orientation=\"vertical\",\n",
    "    fontsize={\n",
    "        \"xtick_label\": GLOBAL_FS,\n",
    "        \"ytick_label\": GLOBAL_FS,\n",
    "        \"cbar_label\": GLOBAL_FS,\n",
    "        \"xlabel\": GLOBAL_FS,\n",
    "        \"ylabel\": GLOBAL_FS,\n",
    "        \"cbar_tick_label\": GLOBAL_FS,\n",
    "        \"title\": GLOBAL_FS,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Add markers for the centres of small holes in the rotated coordinate system\n",
    "for lat, lon in zip(lats_centers_small_holes, longs_centers_small_holes):\n",
    "    newprojplot(\n",
    "        theta=c[0] - lat, phi=c[1] + lon, marker=\"x\", color=\"white\", markersize=5, linewidth=0\n",
    "    )\n",
    "\n",
    "# Draw a rotated grid of latitude lines in the observed patch\n",
    "ndots = 1000\n",
    "latitudes = np.ones(ndots) * np.pi / 2 - np.pi / 3\n",
    "longitudes = np.linspace(\n",
    "    np.pi / 4 + rotation[0] * np.pi / 180, -np.pi / 4 + rotation[0] * np.pi / 180, ndots\n",
    ")\n",
    "rotated_latitudes, rotated_longitudes = hp.Rotator(rot=rotation)(latitudes, longitudes)\n",
    "longitudes_ext = np.linspace(\n",
    "    rotation[0] * np.pi / 180 - np.pi * 0.999, rotation[0] * np.pi / 180 + np.pi * 0.999, ndots\n",
    ")\n",
    "\n",
    "# Visualise latitude bands\n",
    "Nlines = int(1.5 * maxlat * 180 / np.pi)\n",
    "for i in range(Nlines):\n",
    "    latitudes_ext = np.ones(ndots) * np.pi / 2 - maxlat * i / (Nlines - 1)\n",
    "\n",
    "    rotated_latitudes_ext, rotated_longitudes_ext = hp.Rotator(rot=rotation)(\n",
    "        latitudes_ext, longitudes_ext\n",
    "    )\n",
    "    newprojplot(\n",
    "        rotated_latitudes_ext,\n",
    "        rotated_longitudes_ext,\n",
    "        color=\"white\",\n",
    "        alpha=1 - (i / (Nlines - 1)),\n",
    "        linestyle=\"-\",\n",
    "        linewidth=5,\n",
    "    )\n",
    "\n",
    "color_octant = \"orchid\"\n",
    "\n",
    "\n",
    "# Add visual elements for the observed octant and real galactic coordinate grid\n",
    "newprojplot(\n",
    "    theta=line1_theta, phi=line1_phi, color=\"orchid\", markersize=10, label=\"Observed octant\"\n",
    ")\n",
    "newprojplot(theta=line2_theta, phi=line2_phi, color=\"orchid\", markersize=10)\n",
    "newprojplot(theta=line3_theta, phi=line3_phi, color=\"orchid\", markersize=10)\n",
    "newprojplot(\n",
    "    rotated_latitudes,\n",
    "    rotated_longitudes,\n",
    "    color=\"orchid\",\n",
    "    linestyle=\"-.\",\n",
    "    linewidth=2,\n",
    "    markersize=1,\n",
    ")\n",
    "\n",
    "\n",
    "col_real_galactic = \"black\"\n",
    "linestyle_real_galactic = \"--\"\n",
    "\n",
    "lats = np.linspace(-np.pi * 0.999, np.pi * 0.999, ndots)\n",
    "ones = np.ones(ndots)\n",
    "rot_lats, rot_longs = hp.Rotator(rot=rotation)(lats, 0 * ones)\n",
    "newprojplot(\n",
    "    rot_lats, rot_longs, color=col_real_galactic, linestyle=linestyle_real_galactic, linewidth=0.5\n",
    ")\n",
    "for long in [np.pi / 4, np.pi / 2, 3 * np.pi / 4]:\n",
    "    longs = ones * long\n",
    "    rot_lats, rot_longs = hp.Rotator(rot=rotation)(lats, longs)\n",
    "    newprojplot(\n",
    "        rot_lats,\n",
    "        rot_longs,\n",
    "        color=col_real_galactic,\n",
    "        linestyle=linestyle_real_galactic,\n",
    "        linewidth=0.2,\n",
    "    )\n",
    "\n",
    "longs = np.linspace(\n",
    "    rotation[0] * np.pi / 180 - np.pi * 0.999, rotation[0] * np.pi / 180 + np.pi * 0.999, ndots\n",
    ")\n",
    "lats = np.ones(ndots) * np.pi / 2\n",
    "rot_lats, rot_longs = hp.Rotator(rot=rotation)(lats, longs)\n",
    "newprojplot(\n",
    "    rot_lats, rot_longs, color=col_real_galactic, linestyle=linestyle_real_galactic, linewidth=0.5\n",
    ")\n",
    "for lat in [np.pi / 4, 3 * np.pi / 4]:\n",
    "    lats = ones * lat\n",
    "    rot_lats, rot_longs = hp.Rotator(rot=rotation)(lats, longs)\n",
    "    newprojplot(\n",
    "        rot_lats,\n",
    "        rot_longs,\n",
    "        color=col_real_galactic,\n",
    "        linestyle=linestyle_real_galactic,\n",
    "        linewidth=0.2,\n",
    "    )\n",
    "\n",
    "newprojplot(\n",
    "    theta=rot_lat_origin_galactic,\n",
    "    phi=rot_lon_origin_galactic,\n",
    "    marker=\"x\",\n",
    "    color=col_real_galactic,\n",
    "    markersize=5,\n",
    ")\n",
    "\n",
    "xx_ticks_row0 = [-2.76, -1.95, -1.23, -0.42, 0.38, 1.26, 1.88, 2.66]\n",
    "yy_ticks_row0 = [2.92, 3.19, 3.53, 3.8, 3.77, 3.46, 3.12, 2.9]\n",
    "ticks_labs_row0 = [r\"$225$\", r\"$180$\", r\"$135$\", r\"$90$\", r\"$45$\", r\"$0$\", r\"$315$\", r\"$270$\"]\n",
    "for i in range(len(yy_ticks_row0)):\n",
    "    xtick = xx_ticks_row0[i]\n",
    "    ytick = yy_ticks_row0[i]\n",
    "    tick_lab = ticks_labs_row0[i]\n",
    "    plt.text(xtick, ytick, tick_lab, color=\"black\", fontsize=GLOBAL_FS_TINY)\n",
    "\n",
    "xx_ticks_col0 = [0.81, 1.39]\n",
    "yy_ticks_col0 = [2.75, 4.5]\n",
    "ticks_labs_col0 = [r\"$-45$\", r\"$+45$\"]\n",
    "for i in range(len(yy_ticks_col0)):\n",
    "    xtick = xx_ticks_col0[i]\n",
    "    ytick = yy_ticks_col0[i]\n",
    "    tick_lab = ticks_labs_col0[i]\n",
    "    plt.text(xtick, ytick, tick_lab, color=\"black\", fontsize=GLOBAL_FS_TINY)\n",
    "\n",
    "plt.savefig(wd / \"expl/projected_mask_horizontal.png\", dpi=300, bbox_inches=\"tight\", transparent=True)\n",
    "plt.savefig(wd / \"expl/projected_mask_horizontal.pdf\", dpi=300, bbox_inches=\"tight\")\n",
    "plt.show()\n",
    "warnings.filterwarnings(\"default\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Custom aperture size based on grid resolution\n",
    "aperture = 2e-3 if N >= 512 else 4.5e-3 if N >= 256 else 1e-2 if N >= 128 else 5e-2\n",
    "\n",
    "# Filename for the mask with holes\n",
    "name_holes = \"holes{}_{}_ext{}_mask_{}_N{}.npy\".format(\n",
    "    Nholes, str(round(aperture, 5)).replace(\".\", \"\"), round(maxlat * 180 / np.pi), today, N\n",
    ")\n",
    "\n",
    "# Generate or load the mask with holes\n",
    "if not exists(wd / name_holes) or force:\n",
    "    # Rotate the map to align with the observed patch\n",
    "    rotator = hp.rotator.Rotator(rot=rotation, coord=[\"C\", \"G\"], deg=True)\n",
    "    hpx_rotated = rotator.rotate_map_pixel(hpx)\n",
    "    lats = c[0] - grid_sph_latlon[:, :, :, 0]\n",
    "    longs = c[1] + grid_sph_latlon[:, :, :, 1]\n",
    "    grid_mask_holes = hp.get_interp_val(hpx_rotated, lats, longs, lonlat=False)\n",
    "\n",
    "    # Create holes by setting values within the aperture to zero\n",
    "    for lat_hole, long_hole in zip(lats_centers_small_holes, longs_centers_small_holes):\n",
    "        distances = np.sqrt((lats - lat_hole) ** 2 + (longs - long_hole) ** 2)\n",
    "        grid_mask_holes[distances <= aperture] = 0\n",
    "\n",
    "    # Rotate latitudes and longitudes back to galactic coordinates\n",
    "    lats_g_1d, longs_g_1d = inverse_rotation(lats.flatten(), longs.flatten())\n",
    "    lats_g = lats_g_1d.reshape((N, N, N))\n",
    "    longs_g = longs_g_1d.reshape((N, N, N))\n",
    "\n",
    "    # Apply extinction near the Galactic plane\n",
    "    mask = (lats_g > np.pi / 2 - maxlat) & (lats_g < np.pi / 2)\n",
    "    extinction_coeff = (-lats_g[mask] + np.pi / 2) / maxlat\n",
    "    grid_mask_holes[mask] *= extinction_coeff\n",
    "\n",
    "    np.save(wd / name_holes, grid_mask_holes)\n",
    "else:\n",
    "    grid_mask_holes = np.load(wd / name_holes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50412 characters]",
      "text/plain": [
       "<Figure size 1920x1440 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "slice = grid_mask_holes[:,:,int(7*N//8)]\n",
    "plt.imshow(slice, cmap=\"Blues\") ; plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualise the 3D grid mask using Plotly if enabled\n",
    "if plotly:\n",
    "    if N <= 128:\n",
    "        fig = plotly_3d(grid_mask_holes, size=N, L=L, colormap=\"Blues\")\n",
    "    else:\n",
    "        # Downsample the grid for visualisation at higher resolutions\n",
    "        downsample_factor = N // 128\n",
    "        downsampled_field = grid_mask_holes[\n",
    "            ::downsample_factor, ::downsample_factor, ::downsample_factor\n",
    "        ]\n",
    "        fig = plotly_3d(downsampled_field, L=L, colormap=\"Blues\")\n",
    "\n",
    "    fig.show()\n",
    "    # clear_large_plot(fig)  # Clear the Plotly figure to avoid memory issues"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "simbel",
   "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.13.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
