from_weights.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "071092bb-9b24-4b9e-949f-c20b6e5248ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tensorflow as tf\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "from tensorflow.keras import layers, models\n",
    "from loss import custom_loss\n",
    "\n",
    "from OptimizedDataGenerator import OptimizedDataGenerator\n",
    "from loss import *\n",
    "from models import *\n",
    "\n",
    "pi = 3.14159265359\n",
    "\n",
    "maxval=1e9\n",
    "minval=1e-9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b8ab6a7d-46f2-4998-8831-6ec46a2dea33",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"model\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " input_1 (InputLayer)        [(None, 13, 21, 20)]      0         \n",
      "                                                                 \n",
      " q_separable_conv2d (QSepara  (None, 11, 19, 5)        285       \n",
      " bleConv2D)                                                      \n",
      "                                                                 \n",
      " q_activation (QActivation)  (None, 11, 19, 5)         0         \n",
      "                                                                 \n",
      " q_conv2d (QConv2D)          (None, 11, 19, 5)         30        \n",
      "                                                                 \n",
      " q_activation_1 (QActivation  (None, 11, 19, 5)        0         \n",
      " )                                                               \n",
      "                                                                 \n",
      " average_pooling2d (AverageP  (None, 3, 6, 5)          0         \n",
      " ooling2D)                                                       \n",
      "                                                                 \n",
      " q_activation_2 (QActivation  (None, 3, 6, 5)          0         \n",
      " )                                                               \n",
      "                                                                 \n",
      " flatten (Flatten)           (None, 90)                0         \n",
      "                                                                 \n",
      " q_dense (QDense)            (None, 16)                1456      \n",
      "                                                                 \n",
      " q_activation_3 (QActivation  (None, 16)               0         \n",
      " )                                                               \n",
      "                                                                 \n",
      " q_dense_1 (QDense)          (None, 16)                272       \n",
      "                                                                 \n",
      " q_activation_4 (QActivation  (None, 16)               0         \n",
      " )                                                               \n",
      "                                                                 \n",
      " q_dense_2 (QDense)          (None, 14)                238       \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 2,281\n",
      "Trainable params: 2,281\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "# compiles model\n",
    "model=CreateModel((13,21,20),n_filters=5,pool_size=3)\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "33d1434b-5ee1-4f0b-85a7-f516efed08a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.compile(\n",
    "    optimizer=tf.keras.optimizers.Nadam(learning_rate=1e-3),\n",
    "    loss=custom_loss\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "723ffc3d-b455-4a8b-acb1-8dbf3a35ca17",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: './weights-50x12P5x100_bnorm1-checkpoints/training_log.csv'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[7], line 5\u001b[0m\n\u001b[1;32m      3\u001b[0m latest_checkpoint \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmax\u001b[39m(checkpoint_files, key\u001b[38;5;241m=\u001b[39mos\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mgetmtime)\n\u001b[1;32m      4\u001b[0m training_cp_path \u001b[38;5;241m=\u001b[39m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mjoin(base_dir, \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtraining_log.csv\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m----> 5\u001b[0m training_history \u001b[38;5;241m=\u001b[39m \u001b[43mpd\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mread_csv\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtraining_cp_path\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/pandas/io/parsers/readers.py:912\u001b[0m, in \u001b[0;36mread_csv\u001b[0;34m(filepath_or_buffer, sep, delimiter, header, names, index_col, usecols, dtype, engine, converters, true_values, false_values, skipinitialspace, skiprows, skipfooter, nrows, na_values, keep_default_na, na_filter, verbose, skip_blank_lines, parse_dates, infer_datetime_format, keep_date_col, date_parser, date_format, dayfirst, cache_dates, iterator, chunksize, compression, thousands, decimal, lineterminator, quotechar, quoting, doublequote, escapechar, comment, encoding, encoding_errors, dialect, on_bad_lines, delim_whitespace, low_memory, memory_map, float_precision, storage_options, dtype_backend)\u001b[0m\n\u001b[1;32m    899\u001b[0m kwds_defaults \u001b[38;5;241m=\u001b[39m _refine_defaults_read(\n\u001b[1;32m    900\u001b[0m     dialect,\n\u001b[1;32m    901\u001b[0m     delimiter,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    908\u001b[0m     dtype_backend\u001b[38;5;241m=\u001b[39mdtype_backend,\n\u001b[1;32m    909\u001b[0m )\n\u001b[1;32m    910\u001b[0m kwds\u001b[38;5;241m.\u001b[39mupdate(kwds_defaults)\n\u001b[0;32m--> 912\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_read\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilepath_or_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/pandas/io/parsers/readers.py:577\u001b[0m, in \u001b[0;36m_read\u001b[0;34m(filepath_or_buffer, kwds)\u001b[0m\n\u001b[1;32m    574\u001b[0m _validate_names(kwds\u001b[38;5;241m.\u001b[39mget(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnames\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m))\n\u001b[1;32m    576\u001b[0m \u001b[38;5;66;03m# Create the parser.\u001b[39;00m\n\u001b[0;32m--> 577\u001b[0m parser \u001b[38;5;241m=\u001b[39m \u001b[43mTextFileReader\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfilepath_or_buffer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwds\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    579\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m chunksize \u001b[38;5;129;01mor\u001b[39;00m iterator:\n\u001b[1;32m    580\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m parser\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/pandas/io/parsers/readers.py:1407\u001b[0m, in \u001b[0;36mTextFileReader.__init__\u001b[0;34m(self, f, engine, **kwds)\u001b[0m\n\u001b[1;32m   1404\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moptions[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhas_index_names\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m kwds[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mhas_index_names\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n\u001b[1;32m   1406\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles: IOHandles \u001b[38;5;241m|\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 1407\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_engine \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_make_engine\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mengine\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/pandas/io/parsers/readers.py:1661\u001b[0m, in \u001b[0;36mTextFileReader._make_engine\u001b[0;34m(self, f, engine)\u001b[0m\n\u001b[1;32m   1659\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m mode:\n\u001b[1;32m   1660\u001b[0m         mode \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m-> 1661\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles \u001b[38;5;241m=\u001b[39m \u001b[43mget_handle\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1662\u001b[0m \u001b[43m    \u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1663\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1664\u001b[0m \u001b[43m    \u001b[49m\u001b[43mencoding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mencoding\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1665\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcompression\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcompression\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1666\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmemory_map\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mmemory_map\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mFalse\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1667\u001b[0m \u001b[43m    \u001b[49m\u001b[43mis_text\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mis_text\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1668\u001b[0m \u001b[43m    \u001b[49m\u001b[43merrors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mencoding_errors\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstrict\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1669\u001b[0m \u001b[43m    \u001b[49m\u001b[43mstorage_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43moptions\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mstorage_options\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1670\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1671\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m   1672\u001b[0m f \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mhandles\u001b[38;5;241m.\u001b[39mhandle\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/pandas/io/common.py:859\u001b[0m, in \u001b[0;36mget_handle\u001b[0;34m(path_or_buf, mode, encoding, compression, memory_map, is_text, errors, storage_options)\u001b[0m\n\u001b[1;32m    854\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(handle, \u001b[38;5;28mstr\u001b[39m):\n\u001b[1;32m    855\u001b[0m     \u001b[38;5;66;03m# Check whether the filename is to be opened in binary mode.\u001b[39;00m\n\u001b[1;32m    856\u001b[0m     \u001b[38;5;66;03m# Binary mode does not support 'encoding' and 'newline'.\u001b[39;00m\n\u001b[1;32m    857\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m ioargs\u001b[38;5;241m.\u001b[39mencoding \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mb\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;129;01min\u001b[39;00m ioargs\u001b[38;5;241m.\u001b[39mmode:\n\u001b[1;32m    858\u001b[0m         \u001b[38;5;66;03m# Encoding\u001b[39;00m\n\u001b[0;32m--> 859\u001b[0m         handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mopen\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[1;32m    860\u001b[0m \u001b[43m            \u001b[49m\u001b[43mhandle\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    861\u001b[0m \u001b[43m            \u001b[49m\u001b[43mioargs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    862\u001b[0m \u001b[43m            \u001b[49m\u001b[43mencoding\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mioargs\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mencoding\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    863\u001b[0m \u001b[43m            \u001b[49m\u001b[43merrors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43merrors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    864\u001b[0m \u001b[43m            \u001b[49m\u001b[43mnewline\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m    865\u001b[0m \u001b[43m        \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    866\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    867\u001b[0m         \u001b[38;5;66;03m# Binary mode\u001b[39;00m\n\u001b[1;32m    868\u001b[0m         handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mopen\u001b[39m(handle, ioargs\u001b[38;5;241m.\u001b[39mmode)\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: './weights-50x12P5x100_bnorm1-checkpoints/training_log.csv'"
     ]
    }
   ],
   "source": [
    "base_dir = f'./weights-50x12P5x100_bnorm1-checkpoints'\n",
    "checkpoint_files = [os.path.join(base_dir, f) for f in os.listdir(base_dir) if f.endswith('.hdf5')]\n",
    "latest_checkpoint = max(checkpoint_files, key=os.path.getmtime)\n",
    "training_cp_path = os.path.join(base_dir, 'training_log.csv')\n",
    "training_history = pd.read_csv(training_cp_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd6f28cd-b379-4667-9117-7e9583d9ff94",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(f\"Loading model from {latest_checkpoint}\")\n",
    "model.load_weights(latest_checkpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6839572d-8328-4363-8223-1b61e8c0f8ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(training_history['epoch'], training_history['loss'])\n",
    "plt.scatter(training_history['epoch'], training_history['val_loss'])\n",
    "plt.legend(['training', 'validation'])\n",
    "plt.grid(True)\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('NLL loss')\n",
    "plt.tight_layout() \n",
    "\n",
    "plt.savefig(os.path.join(base_dir,'training_hist.png'))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b31ec23-5a53-4ace-a693-11ebdfa4dfd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f235921-4d53-44a3-941c-d4a193b1460c",
   "metadata": {},
   "outputs": [],
   "source": [
    "tfrecords_dir_val = \"/depot/cms/users/das214/tfrecords_2t_val_d8\"\n",
    "test_generator = DG.OptimizedDataGenerator(\n",
    "    load_from_tfrecords_dir = tfrecords_dir_val,\n",
    "    quantize = True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6446272-0a29-4e51-ac93-c47c113c0bad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# predicts test data\n",
    "p_test = model.predict(test_generator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e15a37e-7119-445d-b1bc-066ff8333265",
   "metadata": {},
   "outputs": [],
   "source": [
    "complete_truth = None\n",
    "for _, y in tqdm(test_generator):\n",
    "        if complete_truth is None:\n",
    "            complete_truth = y\n",
    "        else:\n",
    "            complete_truth = np.concatenate((complete_truth, y), axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a9b2df5-aa17-4ebd-b1a1-1f457f6b8f33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# creates df with all predicted values and matrix elements - 4 predictions, all 10 unique matrix elements\n",
    "df = pd.DataFrame(p_test,columns=['x','M11','y','M22','cotA','M33','cotB','M44','M21','M31','M32','M41','M42','M43'])\n",
    "\n",
    "# stores all true values in same matrix as xtrue, ytrue, etc.\n",
    "df['xtrue'] = complete_truth[:,0]\n",
    "df['ytrue'] = complete_truth[:,1]\n",
    "df['cotAtrue'] = complete_truth[:,2]\n",
    "df['cotBtrue'] = complete_truth[:,3]\n",
    "df['M11'] = minval+tf.math.maximum(df['M11'], 0)\n",
    "df['M22'] = minval+tf.math.maximum(df['M22'], 0)\n",
    "df['M33'] = minval+tf.math.maximum(df['M33'], 0)\n",
    "df['M44'] = minval+tf.math.maximum(df['M44'], 0)\n",
    "\n",
    "df['sigmax'] = abs(df['M11'])\n",
    "df['sigmay'] = np.sqrt(df['M21']**2 + df['M22']**2)\n",
    "df['sigmacotA'] = np.sqrt(df['M31']**2+df['M32']**2+df['M33']**2)\n",
    "df['sigmacotB'] = np.sqrt(df['M41']**2+df['M42']**2+df['M43']**2+df['M44']**2)\n",
    "\n",
    "# calculates residuals for x, y, cotA, cotB\n",
    "residuals = df['xtrue'] - df['x']\n",
    "residualsy = df['ytrue'] - df['y']\n",
    "residualsA = df['cotAtrue'] - df['cotA']\n",
    "residualsB = df['cotBtrue'] - df['cotB']\n",
    "\n",
    "# stores results as csv\n",
    "df.to_csv(\"test_4x4.csv\",header=True,index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c322562b-2498-48d4-ba47-39ef59af4b2f",
   "metadata": {},
   "source": [
    "# residuals (mean, stdv)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d6294f0-1962-454d-90d5-c0f3f5c8f923",
   "metadata": {},
   "outputs": [],
   "source": [
    "# x\n",
    "(np.mean(residuals)*75,np.std(residuals)*75)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14339f4e-d28f-4e5b-84fc-8d537c2a62cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# y\n",
    "(np.mean(residualsy)*18.75,np.std(residualsy)*18.75)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5d49f1f-6069-4132-9e10-e6075bedd813",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cotA\n",
    "(np.mean(residualsA)*8,np.std(residualsA)*8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e548c523-abbb-4f66-9262-0d72138cbf98",
   "metadata": {},
   "outputs": [],
   "source": [
    "# cotB\n",
    "(np.mean(residualsB)*0.5,np.std(residualsB)*0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba25f773-bdd4-4e84-a41c-0cd1b04567e8",
   "metadata": {},
   "source": [
    "# x residuals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12b412a2-5b8b-4607-b247-5aa97c3db8af",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(p_test[:,0],bins=np.linspace(-2,2,50))\n",
    "plt.yscale('log')\n",
    "plt.xlabel(r'$x$ [um]')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'x_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1154ff4-fbc8-4a87-935d-73632003fe17",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(residuals,bins=np.linspace(-1,1,50))\n",
    "#plt.yscale('Log')\n",
    "plt.xlabel(r'$x-\\mu_x$ [um]')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'x_res_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ee4ca11-5617-4869-92e3-78996fc210ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = sns.regplot(x=df['xtrue'], y=(df['xtrue']-df['x']), x_bins=np.linspace(-1,1,50), fit_reg=None, marker='.')\n",
    "plt.xlabel(r'True $x$ [um]')\n",
    "plt.ylabel(r'$x-\\mu_x$ [um]')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'x_vs_res.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dddd798-8abf-4e23-a653-a5f435014d6b",
   "metadata": {},
   "source": [
    "# y residuals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52594208-a531-4465-a073-5ad397654903",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(p_test[:,2],bins=np.linspace(-1.5,1.5,50))\n",
    "plt.yscale('log')\n",
    "plt.xlabel(r'$y$ [um]')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'y_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfb95424-da91-4c36-8ed5-96e051aee63d",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(residualsy,bins=np.linspace(-1,1,50))\n",
    "#plt.yscale('Log')\n",
    "plt.xlabel(r'$y-\\mu_y$ [um]')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'y_res_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1759a415-f1a8-4906-a8a2-4efeb32fe76a",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = sns.regplot(x=df['ytrue'], y=(df['ytrue']-df['y']), x_bins=np.linspace(-1,1,50), fit_reg=None, marker='.')\n",
    "plt.xlabel(r'True $y$ [um]')\n",
    "plt.ylabel(r'$y-\\mu_y$ [um]')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'x_vs_res.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a62549da-f856-47d0-b938-fbe970e0729b",
   "metadata": {},
   "source": [
    "# cotA residuals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00e7abfd-a3c8-4242-a427-39e3124670df",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(p_test[:,4],bins=np.linspace(-1.25,1.25,50))\n",
    "plt.yscale('log')\n",
    "plt.xlabel(r'$cotA$')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'cotA_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfc386f9-5672-4d71-af14-7a511ae0fdde",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(residualsA,bins=np.linspace(-0.5,0.5,50))\n",
    "#plt.yscale('Log')\n",
    "plt.xlabel(r'$cotA-\\mu_{cotA}$')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'cotA_res_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1de84f9-162a-4008-b469-49f0ac4874e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = sns.regplot(x=df['cotAtrue']*8, y=(df['cotAtrue']-df['cotA'])*8, x_bins=np.linspace(-8,8,50), fit_reg=None, marker='.')\n",
    "plt.xlabel(r'True $cotA$ [um]')\n",
    "plt.ylabel(r'$cotA-\\mu_{cotA}$ [um]')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'cotA_vs_res.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e71ded81-ac11-40fe-89a9-06b8d324f112",
   "metadata": {},
   "source": [
    "# cotB residuals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "812a9798-ff52-4d59-ac84-ac249d873251",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(p_test[:,6]*0.5,bins=np.linspace(-1,1,50))\n",
    "#plt.yscale('Log')\n",
    "plt.xlabel(r'$cotB$')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'cotB_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "974fa306-a997-46e7-a193-6c5117bb29d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(residualsB*0.5,bins=np.linspace(-1,1,50))\n",
    "#plt.yscale('Log')\n",
    "plt.xlabel(r'$cotB-\\mu_{cotB}$')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'cotB_res_dist.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "793f8563-f335-4f18-acce-2933802d73fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = sns.regplot(x=df['cotBtrue']*0.5, y=(df['cotBtrue']-df['cotB'])*0.5, x_bins=np.linspace(-1,1,100), fit_reg=None, marker='.')\n",
    "plt.xlabel(r'True $cotB$')\n",
    "plt.ylabel(r'$cotB-\\mu_{cotB} $')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'cotB_vs_res.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e253dafe-01be-4e4d-9320-37b5a5a3e84b",
   "metadata": {},
   "source": [
    "# pull"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ab64af5-4f3d-4e6f-85e2-9d883a15c26d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df['pullx'] = (df['xtrue']-df['x'])/df['sigmax']\n",
    "df['pully'] = (df['ytrue']-df['y'])/df['sigmay']\n",
    "df['pullcotA'] = (df['cotAtrue']-df['cotA'])/df['sigmacotA']\n",
    "df['pullcotB'] = (df['cotBtrue']-df['cotB'])/df['sigmacotB']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ee02cae-ff39-49d7-8b72-2a97e219425b",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(df['sigmax'] * 75,bins=np.linspace(0,35,50))\n",
    "plt.xlabel('$\\sigma_{x}$')\n",
    "plt.yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "751ae6fe-e530-48ff-a94b-f1fe1b08a6dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(df['sigmay']* 18.75,bins=np.linspace(0,6,50))\n",
    "plt.xlabel('$\\sigma_{y}$')\n",
    "plt.yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ced64585-c6da-4512-bd4c-2da7e4c4efd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(df['sigmacotA'] * 8,bins=np.linspace(0,1.0,50))\n",
    "plt.xlabel('$\\sigma_{cotA}$')\n",
    "plt.yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea543092-f60a-46f5-8bb3-b38793042e50",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(df['sigmacotB'] * 0.5 ,bins=np.linspace(0,0.25,50))\n",
    "plt.xlabel('$\\sigma_{cotB}$')\n",
    "plt.yscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbe180a1-90f2-4f36-a7ad-b8f91574df8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import curve_fit\n",
    "\n",
    "def gauss(x, A, mu, sigma):\n",
    "    return A * np.exp(-(x - mu) ** 2 / (2 * sigma ** 2))\n",
    "\n",
    "def pull_plot(ax, var, name):\n",
    "    \n",
    "    h = ax.hist(df[var],bins=np.linspace(-5,5,50),histtype='step')\n",
    "    ax.set_xlabel(name)\n",
    "    ax.set_yscale('log')\n",
    "\n",
    "    ydata = h[0]\n",
    "    xdata = h[1][:-1]+3/50.\n",
    "\n",
    "    pars, cov = curve_fit(gauss,xdata,ydata)\n",
    "\n",
    "    xbins = np.linspace(-5,5,100)\n",
    "    ax.plot(xbins,gauss(xbins,pars[0],pars[1],pars[2]),color='black')\n",
    "    ax.set_ylim(0.5,100000)\n",
    "\n",
    "    print('Mean',pars[1])\n",
    "    print('Sigma',pars[2])\n",
    "    \n",
    "    ax.text(-5,2000,\"$\\mu$=\"+str(round(pars[1],2)))\n",
    "    ax.text(-5,1000,\"$\\sigma$=\"+str(round(abs(pars[2]),2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "591a1425-1023-4151-82a0-9266a9523733",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2,2,sharex=True,sharey=True,figsize=(8,6))\n",
    "pull_plot(axes[0][0],'pullx',r'$x$ pull')\n",
    "pull_plot(axes[0][1],'pully',r'$y$ pull')\n",
    "pull_plot(axes[1][0],'pullcotA',r'$\\cot\\alpha$ pull')\n",
    "pull_plot(axes[1][1],'pullcotB',r'$\\cot\\beta$ pull')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'Pull.png')\n",
    "plt.savefig(save_fig_path)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5041b224-f5d5-4794-9e2a-b1a7c9a631df",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34663317-f568-4ac7-89e2-97c9464ee5c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def residual_plot(ax, thisdf, var1, var2, name, scaling=1.0):\n",
    "    \n",
    "    nbins = 15\n",
    "    \n",
    "    var1_scaled = thisdf[var1] * scaling\n",
    "    var2_scaled = thisdf[var2] * scaling\n",
    "    residual_scaled = var1_scaled - var2_scaled\n",
    "    \n",
    "    xmin = np.min(var1_scaled)\n",
    "    xmax = np.max(var1_scaled)\n",
    "    \n",
    "    step = 1.0*(xmax-xmin)/nbins\n",
    "    \n",
    "    x = sns.regplot(x=var1_scaled, y=residual_scaled, x_bins=np.linspace(xmin,xmax,nbins), fit_reg=None, marker='.', ax=ax)\n",
    "    ax.set_xlabel('True ' + name)\n",
    "    ax.set_ylabel('True - predicted ' + name)\n",
    "    \n",
    "    thisdf['residual'+var2] = residual_scaled\n",
    "    print(var1)\n",
    "    \n",
    "    means = []\n",
    "    upbar = []\n",
    "    downbar = []\n",
    "    for i in range(0,nbins):\n",
    "        means += [np.mean(thisdf['residual'+var2][(var1_scaled>xmin + i*step) & (var1_scaled<xmin + (i+1)*step)])]\n",
    "        upbar += [means[i] + np.mean(thisdf['sigma'+var2][(var1_scaled>xmin + i*step) & (var1_scaled<xmin + (i+1)*step)] * scaling)]\n",
    "        downbar += [means[i] - np.mean(thisdf['sigma'+var2][(var1_scaled>xmin + i*step) & (var1_scaled<xmin + (i+1)*step)] * scaling)]\n",
    "    ax.fill_between(x=np.linspace(xmin,xmax,nbins),y1=upbar,y2=downbar, alpha=0.2)\n",
    "\n",
    "def inverse_cot(cota):\n",
    "    a = np.arctan(1.0/cota)\n",
    "    a[np.where(a<0)] = a[np.where(a<0)] + pi\n",
    "    return a    \n",
    "\n",
    "def residual_plot_deg(ax, thisdf, var1, var2, name, scaling=1.0):\n",
    "    # positions\n",
    "    if 'cot' not in var1:\n",
    "        residual_plot(ax, thisdf, var1, var2, name, scaling=scaling)\n",
    "        return\n",
    "\n",
    "    thisdf['angle'] = inverse_cot(thisdf[var2].values * scaling)*180/pi\n",
    "    \n",
    "    thisdf['angleup'] = abs(inverse_cot((thisdf[var2].values + thisdf['sigma'+var2].values) * scaling)*180/pi - thisdf['angle'])\n",
    "    thisdf['angledown'] = abs(inverse_cot((thisdf[var2].values - thisdf['sigma'+var2].values) * scaling)*180/pi - thisdf['angle'])\n",
    "    thisdf['angletrue'] = inverse_cot(thisdf[var1].values * scaling)*180/pi\n",
    "        \n",
    "    var1 = 'angletrue'\n",
    "    var2 = 'angle'\n",
    "    \n",
    "    nbins = 15\n",
    "    xmin = np.min(thisdf[var1])\n",
    "    xmax = np.max(thisdf[var1])\n",
    "    \n",
    "    step = 1.0*(xmax-xmin)/nbins\n",
    "        \n",
    "    x = sns.regplot(x=thisdf[var1], y=(thisdf[var1]-thisdf[var2]), x_bins=np.linspace(xmin,xmax,nbins), fit_reg=None, marker='.', ax=ax)\n",
    "    ax.set_xlabel('True ' + name)\n",
    "    ax.set_ylabel('True - predicted ' + name)\n",
    "    \n",
    "    thisdf['residual'+var2] = (thisdf[var1]-thisdf[var2])\n",
    "    print(var1)\n",
    "    \n",
    "    means = []    \n",
    "    upbar = []\n",
    "    downbar = []\n",
    "    for i in range(0,nbins):\n",
    "        means += [np.mean(thisdf['residual'+var2][(thisdf[var1]>xmin + i*step) & (thisdf[var1]<xmin + (i+1)*step)])]\n",
    "        upbar += [means[i] + np.mean(thisdf['angleup'][(thisdf[var1]>xmin + i*step) & (thisdf[var1]<xmin + (i+1)*step)])]\n",
    "        downbar += [means[i] - np.mean(thisdf['angledown'][(thisdf[var1]>xmin + i*step) & (thisdf[var1]<xmin + (i+1)*step)])]\n",
    "    #ax.scatter(x=np.linspace(xmin,xmax,nbins),y=means)\n",
    "    ax.fill_between(x=np.linspace(xmin,xmax,nbins),y1=upbar,y2=downbar, alpha=0.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc916749-f900-4dd2-9f56-623778168ae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2,2,figsize=(8,6))\n",
    "fig.tight_layout(pad=4.0)\n",
    "residual_plot(axes[0][0],df,'xtrue','x',r'$x$ [um]', scaling=75.0)\n",
    "axes[0][0].plot([-25,-25],[-10,10],color='gray',linestyle=':')\n",
    "axes[0][0].plot([25,25],[-10,10],color='gray',linestyle=':')\n",
    "residual_plot(axes[0][1],df,'ytrue','y',r'$y$ [um]', scaling=18.75)\n",
    "axes[0][1].plot([-6.25,-6.25],[-2,2],color='gray',linestyle=':')\n",
    "axes[0][1].plot([6.25,6.25],[-2,2],color='gray',linestyle=':')\n",
    "residual_plot_deg(axes[1][0],df,'cotAtrue','cotA',r'$\\alpha$ [deg]', scaling=8.0)\n",
    "axes[1][0].plot([90,90],[-10,10],color='gray',linestyle=':')\n",
    "residual_plot_deg(axes[1][1],df,'cotBtrue','cotB',r'$\\beta$ [deg]', scaling=0.5)\n",
    "axes[1][1].plot([90,90],[-10,10],color='gray',linestyle=':')\n",
    "\n",
    "save_fig_path = os.path.join(base_dir, 'summary.png')\n",
    "plt.savefig(save_fig_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bca2f5b-3484-4a09-947c-0db4ce35b22d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "491575ea-fc6c-4741-8a91-eae2d2cc51b7",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a398be0c-5e8a-4eab-934c-697f1e49bee6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:.conda-myenv]",
   "language": "python",
   "name": "conda-env-.conda-myenv-py"
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}