{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PEC model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start by importing MagmaPEC and MagmaPandas and any other packages you want to use. Here we also import Pandas for importing pressure data. For details on the use of MagmaPandas, please see it's [documentation](https://magmapandas.readthedocs.io/en/latest/)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import MagmaPEC as mpc\n",
    "import MagmaPandas as mp\n",
    "\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import your melt inclusion and olivine data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "melt_file = \"./data/melt.csv\"\n",
    "olivine_file = \"./data/olivine.csv\"\n",
    "\n",
    "melt = mp.read_melt(melt_file, index_col=[\"name\"])\n",
    "olivine = mp.read_olivine(olivine_file, index_col=[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have data for internal inclusion pressures (in bars) you can also import those. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "name\n",
       "PI032-04-01    6527.087712\n",
       "PI032-04-02    6828.021536\n",
       "PI041-02-02    4802.872035\n",
       "PI041-03-01    6958.250380\n",
       "PI041-03-03    6826.596778\n",
       "PI041-05-04    4013.630505\n",
       "PI041-05-06    5269.481653\n",
       "PI041-07-01    4969.100477\n",
       "PI041-07-02    4175.590768\n",
       "PI052-01-02    2388.318094\n",
       "Name: 0, dtype: float64"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pressure_file =\"./data/pressure.csv\"\n",
    "\n",
    "pressure = pd.read_csv(pressure_file, index_col = [\"name\"]).squeeze()\n",
    "\n",
    "pressure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, if you have melt CO<sub>2</sub> (and H<sub>2</sub>O) data, you can use MagmaPandas to calculate volatile saturation pressures. First calculate temperature with some initial pressure and then iteratively calculate pressure and temperature. See the [MagmaPandas documentation](https://magmapandas.readthedocs.io/en/latest/notebooks/melt_basics.html) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saturation pressure...    |█████████████████████████| 100% [10/10] in 0.9s \n",
      "Saturation pressure...    |█████████████████████████| 100% [10/10] in 0.9s \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "name\n",
       "PI032-04-01    6156.914986\n",
       "PI032-04-02    6828.021538\n",
       "PI041-02-02    4802.872035\n",
       "PI041-03-01    6958.250380\n",
       "PI041-03-03    6826.596776\n",
       "PI041-05-04    4013.630505\n",
       "PI041-05-06    5269.481651\n",
       "PI041-07-01    4969.100476\n",
       "PI041-07-02    4175.590771\n",
       "PI052-01-02    2388.318094\n",
       "dtype: float64"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "temperature = melt.temperature(P_bar = 5e3)\n",
    "pressure = melt.volatile_saturation_pressure(T_K=temperature)\n",
    "\n",
    "while True:\n",
    "    temperature = melt.temperature(P_bar=pressure)\n",
    "    pressure_new = melt.volatile_saturation_pressure(T_K=temperature)\n",
    "    dP = (pressure_new - pressure) / pressure_new # calculate percentage change\n",
    "    pressure = pressure_new.copy()\n",
    "    if (dP < 0.01).all():\n",
    "        # break when all pressure have converged within 1% of previous values.\n",
    "        break\n",
    "\n",
    "pressure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we will also use a whole-rock dataset to set up a model for predicting the initial FeO contents of our melt inclusions. For explanation of this model, please follow the [Initial FeO prediction](https://magmapec.readthedocs.io/en/latest/notebooks/FeOi.html#) example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wholerock_file = \"./data/wholerock.csv\"\n",
    "\n",
    "wholerock = mp.read_melt(wholerock_file, index_col=[\"name\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's begin by setting up the model to predict the initial FeO content of our melt inclusions according to the previous [example](https://magmapec.readthedocs.io/en/latest/notebooks/FeOi.html#). Here, the FeO prediction model is based on TiO<sub>2</sub>, Al<sub>2</sub>O<sub>3</sub> and CaO, which we can check by accessing their coefficients:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = wholerock.drop(columns=[\"FeO\"])\n",
    "FeOi_predict = mpc.FeOi_prediction(x=x, FeO=wholerock[\"FeO\"])\n",
    "\n",
    "do_not_use = [\"MnO\", \"P2O5\", \"Cr2O3\", \"total\"]\n",
    "\n",
    "model_fits = FeOi_predict.calculate_model_fits(exclude=do_not_use)\n",
    "FeOi_predict.select_predictors(idx=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FeOi_predict.coefficients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To use this model during PEC correction, we need to store this model as a callable function, which we can do with the *model* attribute:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FeO_model = FeOi_predict.model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's first quickly test the model by predicting FeO contents for our uncorrected melt compositions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FeO_model(melt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This all looks good, so let continue with setting up the PEC correction model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, preview the melt and olivine data to check everything looks ok:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "melt.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "olivine.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make sure that each row in *melt* and *olivine* is a matching pair of melt inclusion and olivine host. Here we use the sample names stored in the indices of both dataframes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(melt.index, \"\\n\\n\", olivine.index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can verify that the indices are the same with the *equals* method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "melt.index.equals(olivine.index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the configuration of MagmaPandas and the PEC model and make changes if you want to (see the [configuration example](https://magmapec.readthedocs.io/en/latest/notebooks/config.html)). You can change MagmaPandas settings either via the MagmaPandas object (here: mp.configuration) or the MagmaPEC object (here: mpc.configuration)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(mpc.configuration)\n",
    "print(mpc.PEC_configuration)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can initialise the PEC model, where we need the following data:\n",
    "\n",
    "- **inclusions**\n",
    "\n",
    "    inclusion major element compositions in oxide wt. % as a MagmaPandas Melt frame.\n",
    "\n",
    "- **olivines**\n",
    "\n",
    "    olivine major element compositions in oxide wt. % as a MagmaPandas Olivine frame.\n",
    "\n",
    "- **P_bar**\n",
    "\n",
    "    pressures in bar at which to run the model. You can use a fixed pressure for all inclusions, e.g. *P_bar=2e3* for 2 kbar, or indicate specific pressures per inclusion. In this example we do the latter, by passing the above imported pandas Series with internal inclusion pressures\n",
    "\n",
    "- **FeO_target**\n",
    "\n",
    "    Estimated initial FeO contents of melt inclusions. You can use a fixed value for all melt inclusions, e.g. FeO = 11, specific values for individual melt inclusions, stored in a pandas Series, or a predictive equation based on melt major element composition. This equation needs to be a callable function that accepts a Pandas DataFrame with melt compositions in oxide wt. % as input and return an array-like object with initial FeO contents per inclusion. In the example above we showed how to set up a function like this using MagmaPEC."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pec_model = mpc.PEC(inclusions=melt, olivines=olivine, P_bar=pressure, FeO_target=FeO_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we simply run the model with the *correct* method. It runs in two stages: the first stage equilibrates Fe and Mg between melt inclusions and olivine hosts through Fe-Mg cation exchange and the second stage corrects melt inclusions back to their initial FeO contents by melting or crystallising olivine.\n",
    "\n",
    "This method returns three objects:\n",
    "\n",
    "- corrected melt compositions as a MagmaPandas Melt frame\n",
    "\n",
    "- PEC extents as a pandas DataFrame\n",
    "    \n",
    "- inclusion entrapment temperatures as a pandas Series"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "melts_corrected, pec, temperatures = pec_model.correct()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The dataframe with pec results has three columns:\n",
    "\n",
    "- equilibration_crystallisation: \n",
    "    \n",
    "    crystallisation extents during the equilibration stage\n",
    "\n",
    "- PE_crystallisation:\n",
    "\n",
    "    crystallisation extents during the crystallisation stage\n",
    "\n",
    "- total_crystallisation:\n",
    "\n",
    "    total amount of post-entrapment crystallisation\n",
    "\n",
    "with all data in percent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Results are also stored in the pec object and can be accessed via the *inclusions* and *olivine_corrected* attributes, while entrapment temperatures can be calculated on the fly with the *temperature* method of the *inclusions* Melt frame (see the [MagmaPandas documentation](https://magmapandas.readthedocs.io/en/latest/code_documentation.html#MagmaPandas.MagmaFrames.Melt.temperature))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pec_model.inclusions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pec_model.olivine_corrected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pec_model.inclusions.temperature(P_bar=pressure)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py310",
   "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": 2
}