{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extending Pandas DataFrame Class\n",
    "\n",
    "In this notebook we show how to extend the Pandas DataFrame class using the IsotopeAbundances class defined in tardis/io/decay.py. Official Pandas documentation is located here: https://pandas.pydata.org/pandas-docs/stable/development/extending.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import astropy.units as u\n",
    "from radioactivedecay import Nuclide, Inventory\n",
    "from radioactivedecay.utils import Z_to_elem, elem_to_Z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Notice that we add the new property 'time_0' to _metadata\n",
    "# and that we overwrite the __init__() function.\n",
    "\n",
    "class IsotopeAbundances(pd.DataFrame):\n",
    "\n",
    "    _metadata = [\"time_0\"]\n",
    "    \n",
    "    def __init__(self, *args, **kwargs):\n",
    "        if 'time_0' in kwargs:\n",
    "            time_0 = kwargs['time_0']\n",
    "            kwargs.pop('time_0')\n",
    "        else:\n",
    "            time_0 = 0 * u.d\n",
    "        super(IsotopeAbundances, self).__init__(*args, **kwargs)\n",
    "        self.time_0 = time_0\n",
    "    \n",
    "    @property\n",
    "    def _constructor(self):\n",
    "        return IsotopeAbundances\n",
    "\n",
    "\n",
    "    def _update_inventory(self):\n",
    "        self.comp_dicts = [dict() for i in range(len(self.columns))]\n",
    "        for (atomic_number, mass_number), abundances in self.iterrows():\n",
    "            nuclear_symbol = f'{Z_to_elem(atomic_number)}{mass_number}'\n",
    "            for i in range(len(self.columns)):\n",
    "                self.comp_dicts[i][nuclear_symbol] = abundances[i]\n",
    "\n",
    "    @classmethod\n",
    "    def from_inventories(cls, inventories):\n",
    "        multi_index_tuples = set([])\n",
    "        for inventory in inventories:\n",
    "            multi_index_tuples.update([cls.id_to_tuple(key)\n",
    "                                       for key in inventory.contents.keys()])\n",
    "\n",
    "        index = pd.MultiIndex.from_tuples(\n",
    "            multi_index_tuples, names=['atomic_number', 'mass_number'])\n",
    "\n",
    "\n",
    "        abundances = pd.DataFrame(data=0.0, index=index, columns=range(len(inventories)))\n",
    "\n",
    "        for i, inventory in enumerate(inventories):\n",
    "            for nuclide, abundance in inventory.masses('g').items():\n",
    "                abundances.loc[cls.id_to_tuple(nuclide), i] = abundance\n",
    "\n",
    "        return cls(abundances)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    @staticmethod\n",
    "    def id_to_tuple(atomic_id):\n",
    "        nuclide = Nuclide(atomic_id)\n",
    "        return nuclide.Z, nuclide.A\n",
    "\n",
    "\n",
    "    def to_inventories(self):\n",
    "        \"\"\"\n",
    "        Convert DataFrame to a list of inventories interpreting the MultiIndex as\n",
    "        atomic_number and mass_number\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "            : ~list\n",
    "            list of radioactivedecay Inventories\n",
    "        :return:\n",
    "        \"\"\"\n",
    "\n",
    "        comp_dicts = [dict() for i in range(len(self.columns))]\n",
    "        for (atomic_number, mass_number), abundances in self.iterrows():\n",
    "            nuclear_symbol = f'{Z_to_elem(atomic_number)}{mass_number}'\n",
    "            for i in range(len(self.columns)):\n",
    "                comp_dicts[i][nuclear_symbol] = abundances[i]\n",
    "        return [Inventory(comp_dict, 'g') for comp_dict in comp_dicts]\n",
    "\n",
    "\n",
    "\n",
    "    def decay(self, t):\n",
    "        \"\"\"\n",
    "        Decay the Model\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "\n",
    "        t: ~float or ~astropy.units.Quantity\n",
    "            if float it will be understood as days\n",
    "\n",
    "        Returns:\n",
    "            : decayed abundances\n",
    "        \"\"\"\n",
    "\n",
    "        inventories = self.to_inventories()\n",
    "        t_second = u.Quantity(t, u.day).to(u.s).value - self.time_0.to(u.s).value\n",
    "        decayed_inventories = [item.decay(t_second, 's') for item in inventories]\n",
    "        df = IsotopeAbundances.from_inventories(decayed_inventories)\n",
    "        df.sort_index(inplace=True)\n",
    "        return df \n",
    "\n",
    "    def as_atoms(self):\n",
    "        \"\"\"\n",
    "        Merge Isotope dataframe according to atomic number \n",
    "\n",
    "        Returns:\n",
    "            : merged isotope abundances\n",
    "        \"\"\"\n",
    "\n",
    "        return self.groupby('atomic_number').sum()\n",
    "\n",
    "    def merge(self, other, normalize=True):\n",
    "        \"\"\"\n",
    "        Merge Isotope dataframe with abundance passed as parameter \n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        other: pd.DataFrame \n",
    "        normalize : bool\n",
    "            If true, resultant dataframe will be normalized\n",
    "\n",
    "        Returns:\n",
    "            : merged abundances\n",
    "        \"\"\"\n",
    "        isotope_abundance = self.as_atoms()\n",
    "        isotope_abundance = isotope_abundance.fillna(0.0)\n",
    "        #Merge abundance and isotope dataframe\n",
    "        modified_df = isotope_abundance.add(other, fill_value=0)\n",
    "\n",
    "        if normalize:\n",
    "            norm_factor = modified_df.sum(axis=0)\n",
    "            modified_df /= norm_factor\n",
    "\n",
    "        return modified_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = IsotopeAbundances({'A': [1, 2, 3], 'B': [4, 5, 6], 'C': [7, 8, 9]}, time_0=5*u.d)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now the new property is accessible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$5 \\; \\mathrm{d}$"
      ],
      "text/plain": [
       "<Quantity 5. d>"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.time_0"
   ]
  }
 ],
 "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}