{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "77780177-041d-44f9-817d-8f9d15796b8c",
   "metadata": {},
   "source": [
    "# Spectral cube modifications (apodization, interpolation etc.)\n",
    "\n",
    "You can use the following functions to modify your cube data and export it to a new hdf5 cube (ORCS compatible) or a fits cube.\n",
    "\n",
    "Even if possible, you may want to use a **level 3** cube since level 3 cubes contains complex data which allows for perfect interpolation, apodization etc. Older cubes contains only the real spectrum.\n",
    "\n",
    "Note that these functions can take a very long time to run :)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1c32a22-d9da-449c-bf53-472558c841bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from orcs.process import SpectralCube\n",
    "import orb.utils.io\n",
    "import logging\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba28e246-d906-4d69-b7a9-0b56085e3913",
   "metadata": {},
   "source": [
    "## Apodization (exported to a new hdf5 cube that can be used with ORCS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be59d7a5-05ec-43f6-8a0d-dc6289ee0e7e",
   "metadata": {},
   "source": [
    "### With a level 2, 2.5 or 3 cube"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8bf1874-5c9f-4c76-ba85-7daf5fadcd79",
   "metadata": {},
   "outputs": [],
   "source": [
    "cube = SpectralCube('M1_2022_SN3.merged.cm1.hdf5')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a70b90f7-fb61-4f44-91b1-418d6322d78f",
   "metadata": {},
   "outputs": [],
   "source": [
    "cube_path = 'NGC6888_North_SN3.merged.cm1.1.0.hdf5'\n",
    "# You must first make a **hardcopy** of the original cube and rename it (here as outcube.hdf5)\n",
    "export_path = 'outcube.hdf5'\n",
    "\n",
    "def apodize_cube(cube_path, export_path, coeff):\n",
    "    cube = SpectralCube(cube_path)\n",
    "    if not np.iscomplex(cube[1000,1000,0]):\n",
    "        print('WARNING: cube is not complex, opt for a level 3 cube')\n",
    "        \n",
    "    logging.getLogger().setLevel(logging.CRITICAL)\n",
    "    fakespec = cube.get_spectrum(0,0, r=0)\n",
    "    with cube.open_hdf5() as f:\n",
    "        with orb.utils.io.open_hdf5(export_path, 'a') as fout:\n",
    "            for ii in range(cube.dimx):\n",
    "                print('\\r', ii, '/', cube.dimx, end='')\n",
    "                for ij in range(cube.dimy):\n",
    "                    ispec = cube[ii,ij,:]\n",
    "                    fakespec.data = ispec\n",
    "                    # here is the apodization function. You can change this function and export it\n",
    "                    # to an ORCS-compatible HDF5 cube as long as you don't change the axis of the data \n",
    "                    # (no interpolation)\n",
    "                    fakespecA = fakespec.apodize(coeff)\n",
    "                    # e.g. if you want to take the absolute value instead of the real part : \n",
    "                    # fakespecA.data = np.abs(fakespecA.data).astype(complex)\n",
    "                    \n",
    "                    # you may now write the new spectrum in place of the old data\n",
    "                    fout['data'][ii,ij,:] = fakespecA.data\n",
    "    logging.getLogger().setLevel(logging.INFO)\n",
    "                    \n",
    "apodize_cube(cube_path, export_path, 2) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74974525-9e68-412f-b8ae-1d643f820d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.getLogger().setLevel(logging.INFO)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6220296-d581-4582-b05b-714aaff825e1",
   "metadata": {},
   "source": [
    "### With a level 1 cube"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a44ea8f-7bf7-4522-af5b-e1d0fdbcda89",
   "metadata": {},
   "outputs": [],
   "source": [
    "cube_path = 'M57_SN3.merged.cm1.1.0.hdf5'\n",
    "# You must first make a **hardcopy** of the original cube and rename it (here as outcube.hdf5)\n",
    "export_path = 'outcube_old.hdf5'\n",
    "\n",
    "cube = SpectralCube(cube_path)\n",
    "\n",
    "def apodize_old_cube(cube_path, export_path, coeff):\n",
    "    cube = SpectralCube(cube_path)\n",
    "    if not np.iscomplex(cube[1000,1000,0]):\n",
    "        print('WARNING: cube is not complex, opt for a level 3 cube')\n",
    "    fakespec = cube.get_spectrum(0,0, r=0)\n",
    "    logging.getLogger().setLevel(logging.CRITICAL)\n",
    "    with cube.open_hdf5() as f:\n",
    "        with orb.utils.io.open_hdf5(export_path, 'a') as fout:\n",
    "            for iquad in range(9):\n",
    "                quadpath = cube.oldcube._get_hdf5_quad_data_path(iquad)\n",
    "                quad = f[quadpath]\n",
    "                quadout = fout[quadpath]\n",
    "                print('quadrant', iquad, '/ 9 =======')\n",
    "                for ii in range(quad.shape[0]):\n",
    "                    print('\\r', ii, '/', cube.shape[0], end='')\n",
    "                    for ij in range(quad.shape[1]):\n",
    "                        ispec = quad[ii,ij,:]\n",
    "                        fakespec.data = ispec\n",
    "                        fakespec = fakespec.apodize(coeff)\n",
    "                        quadout[ii,ij,:] = fakespec.data.real\n",
    "    logging.getLogger().setLevel(logging.INFO)\n",
    "                        \n",
    "apodize_old_cube(cube_path, export_path, 2) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "679de155-1abb-4f04-8d93-8466fab1d8b7",
   "metadata": {},
   "source": [
    "## Interpolation on a new axis (here a loglinear axis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a2bd43d-ea85-4c3e-8b81-52ff2b46bf38",
   "metadata": {},
   "outputs": [],
   "source": [
    "import orb.utils.io\n",
    "import orb.utils.spectrum\n",
    "\n",
    "cube = SpectralCube('NGC6888_North_SN3.merged.cm1.1.0.hdf5')\n",
    "\n",
    "cube_axis = cube.get_axis(1000, 1000)\n",
    "\n",
    "# create the projection axis\n",
    "loglin_axis = 10**np.linspace(np.log10(cube_axis.data[0]), np.log10(cube_axis.data[-1]), 1000)\n",
    "\n",
    "export_path = 'outcube.fits'\n",
    "\n",
    "# get flambda depending on the level of the cube, to export the flux calibration with the data\n",
    "if not np.iscomplex(cube[1000,1000,0]):\n",
    "    print('WARNING: cube is not complex, opt for a level 3 cube')\n",
    "\n",
    "flambda = np.ones(cube.dimz, dtype=float)\n",
    "if cube.get_level() >= 3:\n",
    "    if 'flambda' in cube.params:\n",
    "        flambda = cube.params.flambda\n",
    "    \n",
    "    if np.size(flambda) == 1:\n",
    "        flambda *= np.ones(cube.dimz, dtype=float)\n",
    "    elif np.size(flambda) != cube.dimz:\n",
    "        logging.warning('bad flux calibration, output will not be flux calibrated')\n",
    "        flambda = np.ones(cube.dimz, dtype=float)\n",
    "\n",
    "\n",
    "# write data to a HDF5 file (impossible to do with a FITS file without loading the whole data in memory)\n",
    "with orb.utils.io.open_hdf5('test.hdf5', 'w') as fout:\n",
    "    fout.create_dataset('/data', shape=(cube.dimx, cube.dimy, len(loglin_axis)), dtype=float)\n",
    "    for ii in range(cube.dimx):\n",
    "        print('\\r', ii, '/', cube.dimx, end='')\n",
    "        for ij in range(cube.dimy):\n",
    "            ispec = cube[ii,ij,:]\n",
    "            newspec = orb.utils.spectrum.project(ispec, cube_axis.data, loglin_axis, 10).astype(float)\n",
    "            fout['data'][ii,ij,:] = newspec\n",
    "          \n",
    "        \n",
    "# export the HDF5 to a FITS cube\n",
    "with orb.utils.io.open_hdf5('test.hdf5', 'r') as f:\n",
    "    shdu = orb.utils.io.open_large_fits_file(cube, export_path)\n",
    "    \n",
    "    for iz in range(f['data'].shape[2]):\n",
    "        print('\\r', iz, 'Exporting frame {}'.format(iz), end='')\n",
    "        shdu.write(f['data'][:,:,iz].real.astype(np.float32).T * flambda[iz])\n",
    "        \n",
    "    shdu.close()\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b2f7193-bdc0-43f3-aae4-f1e25cdde1c5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}