{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "graduate-saying",
   "metadata": {},
   "source": [
    "Note on pytorch. NN optimization acts imperitively/by side effect as follows.\n",
    " - Define model\n",
    " - loop\n",
    "     - Calculate loss\n",
    "     - Zero gradients\n",
    "     - backprop to model\n",
    "     - check conditions for exit\n",
    " - report diagnostics\n",
    " - disect results\n",
    " \n",
    " \n",
    "## Split result from NN\n",
    "Goal is to train the NN and then get a couple of outputs at the end that can be used to split between value function partials and launch functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "lasting-portable",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import combined as c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "marked-paris",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class LaunchFnEstimand(torch.nn.Module):\n",
    "    def __init__(self, state_tensor_size,layers_size,number_constellations):\n",
    "        \"\"\"\n",
    "        Description\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.number_constellations = number_constellations\n",
    "        self.layers_size = layers_size\n",
    "        self.state_tensor_size = state_tensor_size\n",
    "        \n",
    "        #Scale up the input from just the tensor of states to the layer_size X number_constellations\n",
    "        \n",
    "        #Increase to the layer size\n",
    "        self.linear_1 = torch.nn.Linear(in_features=state_tensor_size, out_features=layers_size)\n",
    "        self.relu = torch.nn.ReLU()\n",
    "        self.linear_3 = torch.nn.Linear(in_features=layers_size, out_features=layers_size)\n",
    "        self.linear_5 = torch.nn.Linear(in_features=layers_size, out_features=number_constellations)\n",
    "\n",
    "        \n",
    "    def forward(self, input_values):\n",
    "        \n",
    "        intermediate_values = self.relu(input_values) #states should be positive anyway.\n",
    "        intermediate_values = self.linear_1(intermediate_values)\n",
    "        intermediate_values = self.linear_3(intermediate_values)\n",
    "        intermediate_values = self.linear_5(intermediate_values)\n",
    "        intermediate_values = self.relu(intermediate_values) #launches are always positive\n",
    "        \n",
    "        return intermediate_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "artificial-gilbert",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.7175, grad_fn=<AddBackward0>)\n",
      "tensor(0.2107, grad_fn=<AddBackward0>)\n",
      "tensor(0.0724, grad_fn=<AddBackward0>)\n",
      "tensor(0.0259, grad_fn=<AddBackward0>)\n",
      "tensor(0.0094, grad_fn=<AddBackward0>)\n",
      "tensor(0.0034, grad_fn=<AddBackward0>)\n",
      "tensor(0.0012, grad_fn=<AddBackward0>)\n",
      "tensor(0.0004, grad_fn=<AddBackward0>)\n",
      "tensor(0.0002, grad_fn=<AddBackward0>)\n",
      "tensor(5.8468e-05, grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[0.0046, 0.0000]]], grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "launch = LaunchFnEstimand(3,12,2)\n",
    "\n",
    "#Prep Optimizer\n",
    "optimizer = torch.optim.SGD(launch.parameters(),lr=0.01)\n",
    "\n",
    "#get loss function\n",
    "def loss_fn5(output):\n",
    "    return sum(sum(sum((output)**2)))\n",
    "\n",
    "for i in range(10):\n",
    "    #training loop\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    output = launch.forward(test)\n",
    "\n",
    "    l = loss_fn5(output)\n",
    "\n",
    "    l.backward()\n",
    "\n",
    "    optimizer.step()\n",
    "\n",
    "    print(l)\n",
    "    \n",
    "\n",
    "launch.forward(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "political-manchester",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[0.0000, 0.9998]]], grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "launch(test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "delayed-bikini",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EstimandNN(torch.nn.Module):\n",
    "    def __init__(self, state_tensor_size,layers_size,number_constellations):\n",
    "        super().__init__()\n",
    "        \n",
    "        #So, this next section constructs different layers within the NN\n",
    "        #sinlge linear section\n",
    "        \n",
    "        self.partials_estimator = PartialDerivativesEstimand(state_tensor_size,layers_size,number_constellations) #TODO\n",
    "        self.launch_estimator = LaunchFnEstimand(state_tensor_size,layers_size,number_constellations)\n",
    "        \n",
    "    def forward(self, input_values):\n",
    "        partials = self.partials_estimator(input_values)\n",
    "        launch = self.launch_estimator(input_values)\n",
    "        \n",
    "        return c.EstimandInterface(partials,launch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "stable-edmonton",
   "metadata": {},
   "outputs": [],
   "source": [
    "enn = EstimandNN(3,12,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "breeding-afghanistan",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Launch Decisions and Partial Derivativs of value function with\n",
      "\tlaunches\n",
      "\t\t tensor([[[0.0000, 0.0020]]], grad_fn=<ReluBackward0>)\n",
      "\tPartials\n",
      "\t\ttensor([[0.0000, 0.0000],\n",
      "        [1.7938, 1.7938],\n",
      "        [0.0000, 0.0000],\n",
      "        [2.8751, 2.8751],\n",
      "        [1.4894, 1.4894],\n",
      "        [1.4614, 1.4614],\n",
      "        [0.0000, 0.0000],\n",
      "        [2.9800, 2.9800],\n",
      "        [0.0000, 0.0000],\n",
      "        [0.0000, 0.0000],\n",
      "        [0.0000, 0.0000],\n",
      "        [0.0000, 0.0000]], grad_fn=<ReluBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(enn(test))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "applicable-relay",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
