{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c054d421-9717-4feb-bb30-d080d91ceac5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from anisotropy.core.database import Database, Structure\n",
    "from pandas import DataFrame, Series\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn\n",
    "import numpy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7ec2ae1b-8a7e-4d50-b1bd-201fc3cd190d",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = Database(\"anisotropy\", \"woPrismaticLayer\")\n",
    "db.setup()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "75498f51-fa01-494e-a64e-e650138d1ca2",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = db.search([Structure.type == \"simple\", Structure.direction == str([1.0, 0.0, 0.0])])\n",
    "df = DataFrame(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3bc80bca-03ce-416e-be2f-ab7ae2af95fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_numeric = df[[\n",
    "    col for col in df.columns \n",
    "        if not isinstance(df[col][0], str) \n",
    "            and not isinstance(df[col][0], numpy.bool_)\n",
    "            and not isinstance(df[col][0], dict)\n",
    "            and not isinstance(df[col][0], list)\n",
    "            and not df[col][0] is None\n",
    "            and not col[-3: ] == \"_id\"\n",
    "]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "eab92bcc-8f1f-4490-8dc6-8ebf191d6cf2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<AxesSubplot:>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 236064 characters]\n",
      "text/plain": [
       "<Figure size 2160x1440 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "seaborn.set(rc = { \"figure.figsize\": (30, 20) })\n",
    "seaborn.heatmap(df_numeric.corr(), annot = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7611c892-1e13-404e-849e-3cfa15d8dfca",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = df_numeric[[\"theta\", \"r0\", \"L\", \"radius\"]] #.drop(columns = [\"flowRate\"])\n",
    "y = df_numeric[\"flowRate\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1a38aaa0-89fb-40f7-abce-e500ab696349",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "xtr, xte, ytr, yte = train_test_split(x, y, test_size = 0.2, random_state = 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1b7e9f62-35a9-4a1a-9ffa-a2a6e757ff46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn import preprocessing\n",
    "scaler = preprocessing.MinMaxScaler()\n",
    "x_scaled = scaler.fit_transform(xtr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "19bd2f52-a5bf-404c-93f7-0626bc35a915",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DecisionTreeRegressor(random_state=500)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.tree import DecisionTreeRegressor\n",
    "neigh = DecisionTreeRegressor(random_state = 500)\n",
    "neigh.fit(x_scaled, ytr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d063265f-2f06-4290-811e-b6b1753bcac9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.664398739090909e-15"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sklearn.metrics import mean_absolute_error\n",
    "xte_scaled = scaler.transform(xte)\n",
    "y_pred = neigh.predict(xte_scaled)\n",
    "mean_absolute_error(yte, y_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "762af50c-371a-4ba5-b5ab-763970ef2a37",
   "metadata": {},
   "outputs": [],
   "source": [
    "#df_numeric[[\"theta\", \"r0\", \"L\", \"radius\", \"flowRate\", \"volumeCell\", \"volume\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "7a951647-efa1-4231-b96d-f67bc0ed6c33",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>theta</th>\n",
       "      <th>r0</th>\n",
       "      <th>L</th>\n",
       "      <th>radius</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.29</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.408451</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   theta   r0    L    radius\n",
       "0   0.29  1.0  2.0  1.408451"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df = DataFrame([{\n",
    "    \"theta\": 0.29,\n",
    "    \"r0\": 1.,\n",
    "    \"L\": 2.,\n",
    "    \"radius\": 1. / (1. - 0.29)\n",
    "}]); test_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bc9204db-bde4-4032-b4f5-313e99e25ab5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([4.53058768e-15])"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "neigh.predict(scaler.transform(test_df))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b95277f-c49e-42b9-8d5f-3334d5ca7729",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c56035b1-18d5-4c88-94b6-da4d7eccac9a",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
