{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3f3bae16",
   "metadata": {},
   "source": [
    "# 2. Advanced Chemistry (Dynamic Zones)\n",
    "\n",
    "`PA3Py` was designed to be completely composition-agnostic. You can inject your own Python chemical functions and it will automatically discover the species and build the composition arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4e0bd72b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-13T22:23:23.137936Z",
     "iopub.status.busy": "2026-07-13T22:23:23.137615Z",
     "iopub.status.idle": "2026-07-13T22:23:24.589974Z",
     "shell.execute_reply": "2026-07-13T22:23:24.588813Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[load_tripodpy_hdf5] Reading 100 snapshots from ../../tests/test_data/run_smooth_a0.001_v10...\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "import os\n",
    "sys.path.insert(0, os.path.abspath('../../src'))\n",
    "from pa3py import PA3Py\n",
    "\n",
    "# Load simulation\n",
    "sim = PA3Py('../../tests/test_data/run_smooth_a0.001_v10')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42131f45",
   "metadata": {},
   "source": [
    "## Custom Composition Function\n",
    "\n",
    "We will define a protoplanetary disk with 4 chemical species (`silicates`, `H2O`, `CO2`, `CO`) partitioned into 4 dynamic zones."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bf732a8c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-13T22:23:24.592756Z",
     "iopub.status.busy": "2026-07-13T22:23:24.592422Z",
     "iopub.status.idle": "2026-07-13T22:23:24.599405Z",
     "shell.execute_reply": "2026-07-13T22:23:24.598708Z"
    }
   },
   "outputs": [],
   "source": [
    "from pa3py import constants as c\n",
    "\n",
    "def chemistry_4_zones(r_cm, t_sec):\n",
    "    # 1. Define the snowlines (some can migrate)\n",
    "    r_h2o = 2.73 * c.AU * (max(t_sec, 1e-6) / 1e13)**(-0.5)\n",
    "    r_co2 = 5.0 * c.AU\n",
    "    r_co  = 12.0 * c.AU\n",
    "    \n",
    "    # 2. Assign relative abundances based on distance to the star\n",
    "    if r_cm < r_h2o:\n",
    "        return {'silicates': 1.0}                           # 100% rocky\n",
    "    elif r_cm < r_co2:\n",
    "        return {'silicates': 0.5, 'H2O': 0.5}               # Ice Zone\n",
    "    elif r_cm < r_co:\n",
    "        return {'silicates': 0.3, 'H2O': 0.3, 'CO2': 0.4}   # CO2 Zone\n",
    "    else:\n",
    "        return {'silicates': 0.2, 'H2O': 0.2, 'CO2': 0.3, 'CO': 0.3} # Cold Zone\n",
    "\n",
    "# Inject our function into the engine. PA3Py will auto-detect the species.\n",
    "sim.set_custom_chemistry(chemistry_4_zones)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7720d655",
   "metadata": {},
   "source": [
    "## Running Custom Chemistry\n",
    "\n",
    "We will place embryos in various regions and let PA3Py handle the rest."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c93debce",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-07-13T22:23:24.602383Z",
     "iopub.status.busy": "2026-07-13T22:23:24.602165Z",
     "iopub.status.idle": "2026-07-13T22:23:24.627619Z",
     "shell.execute_reply": "2026-07-13T22:23:24.626664Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "---------------------------------------------------------------------------------\n",
      "  r [AU]  M_tot [ME]  M_iso [ME]   f_CO[%]  f_silicates[%]  f_CO2[%]  f_H2O[%]\n",
      "---------------------------------------------------------------------------------\n",
      "    2.00        0.164         6.65       0.0            87.7       0.0      12.3\n",
      "    4.00        0.306        11.19       0.0            63.8       0.0      36.2\n",
      "    8.00        0.123        18.82       0.0            39.7      34.4      25.8\n",
      "   15.00        0.001        30.15       6.5            82.7       6.5       4.3\n",
      "---------------------------------------------------------------------------------\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Embryos in the rocky, ice, CO2, and CO zones\n",
    "results = sim.run_growth([2.0, 4.0, 8.0, 15.0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.14",
   "language": "python",
   "name": "python3.14"
  },
  "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.14.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}