FittingImprovements.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "587a0c62-b36a-4fb7-b223-5625da605abe",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "70c25665-dda0-4950-8690-ff7f954684ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import pyro\n",
    "import pickle as pkl\n",
    "from fitting.high_level import RegressionModel\n",
    "from fitting.regression import DataValues\n",
    "from fitting.high_level import makeDiagnosticPlots\n",
    "import gpytorch\n",
    "import matplotlib.pyplot as plt\n",
    "from fitting.utils import affineTransformMVN\n",
    "import pyro \n",
    "from fitting.high_level import RegressionModel\n",
    "from fitting.regression import DataValues, makeRegressionData\n",
    "from fitting.utils import getScaledEigenvecs, fixMVN, affineTransformMVN\n",
    "from fitting.transformations import getNormalizationTransform\n",
    "import fitting.models as models\n",
    "import fitting.regression as regression\n",
    "import fitting.transformations as transformations\n",
    "import fitting.windowing as windowing\n",
    "from gpytorch.kernels import ScaleKernel as SK\n",
    "import hist\n",
    "import tqdm\n",
    "\n",
    "\n",
    "\n",
    "import pyro.distributions as pyrod\n",
    "import pyro.infer as pyroi\n",
    "\n",
    "\n",
    "from fitting.plot_tools import createSlices, getPolyFromSquares, makeSquares, simpleGrid\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "46136487-745d-4890-8559-2989d4d90ea9",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\n",
    "    \"regression_results/2018_Control_nn_uncomp_0p67_m14_vs_mChiUncompRatio.pkl\",\n",
    "    \"rb\",\n",
    ") as f:\n",
    "    control = pkl.load(f)\n",
    "\n",
    "inhist = control[\"Data2018\", \"Control\"][\"hist_collection\"][\"histogram\"]#[hist.loc(1000) :, hist.loc(0) :]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fbf8c3c8-465b-4b6d-8e25-e1826bbbe553",
   "metadata": {},
   "outputs": [],
   "source": [
    "class LargeFeatureExtractor(torch.nn.Sequential):\n",
    "    def __init__(self):\n",
    "        super(LargeFeatureExtractor, self).__init__()\n",
    "        self.add_module('linear1', torch.nn.Linear(2, 32))\n",
    "        self.add_module('relu1', torch.nn.ReLU())\n",
    "        self.add_module('linear2', torch.nn.Linear(32, 16))\n",
    "        self.add_module('relu2', torch.nn.ReLU())\n",
    "        self.add_module('linear3', torch.nn.Linear(16, 8))\n",
    "        self.add_module('relu3', torch.nn.ReLU())\n",
    "        self.add_module('linear4', torch.nn.Linear(8, 2))\n",
    "\n",
    "ft = LargeFeatureExtractor()\n",
    "class GPRegressionModel(gpytorch.models.ExactGP):\n",
    "        def __init__(self, train_x, train_y, likelihood):\n",
    "            super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)\n",
    "            self.mean_module = gpytorch.means.ConstantMean()\n",
    "            \n",
    "            self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=2))\n",
    "\n",
    "            self.feature_extractor = ft\n",
    "            self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)\n",
    "\n",
    "        def forward(self, x):\n",
    "            projected_x = self.feature_extractor(x)\n",
    "            projected_x = self.scale_to_bounds(projected_x)\n",
    "            mean_x = self.mean_module(projected_x)\n",
    "            covar_x = self.covar_module(projected_x)\n",
    "            return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1523599a-9ecd-4e8c-b837-845df1289ed9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[  -1.6730],\n",
       "        [-999.4132],\n",
       "        [  -3.2343],\n",
       "        [  -3.2343],\n",
       "        [  -3.2343],\n",
       "        [  -3.2343]], grad_fn=<AddmmBackward0>)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class RBFLayer(torch.nn.Module):\n",
    "    def __init__(self, dim, count):\n",
    "        super().__init__()\n",
    "        self.length_scales = torch.nn.Parameter(torch.ones(count))\n",
    "        self.centers = torch.nn.Parameter(torch.ones(count,dim))\n",
    "\n",
    "    def forward(self, vals):\n",
    "        return torch.exp((torch.unsqueeze(vals,1) - self.centers).pow(2).sum(-1)/(2 * self.length_scales))\n",
    "\n",
    "r = RBFLayer(2,3)       \n",
    "a = torch.tensor([[0.9041,  0.0196], [-0.3108, -2.4423], [-0.4821,  1.059], [-0.4821,  1.059], [-0.4821,  1.059], [-0.4821,  1.059]])\n",
    "l =  torch.nn.Linear(3,1)\n",
    "\n",
    "l(r(a))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0a770b6a-0966-47e8-969b-6928ba2feadc",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NonStatKernel(gpytorch.kernels.RBFKernel):\n",
    "    # the sinc kernel is stationary\n",
    "    is_stationary = False\n",
    "    def __init__(self, dim=2, count=4, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.pre_transform = RBFLayer(dim,count)\n",
    "        self.trans = torch.nn.Linear(count,1)\n",
    "        #self.add_module(\"trans\", self.trans)\n",
    "        #self.add_module(\"pre\", self.pre_transform)\n",
    "\n",
    "    # this is the kernel function\n",
    "    def forward(self, x1, x2, diag=False, **params):\n",
    "\n",
    "        v1 = self.trans(self.pre_transform(x1))\n",
    "        v2 = self.trans(self.pre_transform(x2))\n",
    "        r =  super().forward(x1,x2,diag=diag,**params)\n",
    "\n",
    "        if diag:\n",
    "            o = torch.squeeze(v1*v2) \n",
    "        else:\n",
    "            o =torch.outer(v1.squeeze(), v2.squeeze())\n",
    "\n",
    "        return o * r #, **params)\n",
    "        \n",
    "class NonStatGP(gpytorch.models.ExactGP):\n",
    "        def __init__(self, train_x, train_y, likelihood, function=None):\n",
    "            super().__init__(train_x, train_y, likelihood)\n",
    "            self.mean_module = gpytorch.means.ConstantMean()\n",
    "\n",
    "            self.base_covar_module = NonStatKernel(ard_num_dims=2)  #* NonStatKernel(ard_num_dims=2)\n",
    "            if False:\n",
    "                self.covar_module = self.base_covar_module                                                   \n",
    "            else:\n",
    "                self.covar_module = gpytorch.kernels.InducingPointKernel(self.base_covar_module,\n",
    "                            likelihood=likelihood,\n",
    "                             inducing_points=train_x[::2].clone())\n",
    "\n",
    "           # self.covar_module = SK(NonStatKernel(ard_num_dims=2))\n",
    "        def forward(self, x):\n",
    "            mean_x = self.mean_module(x)\n",
    "            covar_x = self.covar_module(x)\n",
    "            return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "630696f1-84c9-45eb-a60a-31ab8e15c94d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def randSample(*args, size=None):\n",
    "    p = torch.randperm(args[0].size(0))\n",
    "    return [x[p[:size]] for x in args]\n",
    "\n",
    "def batchify(*args, size=None, batches=None):\n",
    "    vals = [randSample(*args, size=size) for i in range(batches)]\n",
    "    return tuple(torch.stack(x) for x in zip(*vals))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "025dc91e-a8ae-4430-969d-8b3e35c8c353",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chi2bins(obs,exp,var):\n",
    "    return torch.sum((obs-exp).pow(2) / var ) / obs.size(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "cc7698b3-7ddf-45fa-932d-a6ac68ddc063",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([60, 60])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f67b8332380>"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAaAAAAGfCAYAAAAZGgYhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjkuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy80BEi2AAAACXBIWXMAAA9hAAAPYQGoP6dpAAAivklEQVR4nO3df2yV5f3/8dc5bc9ppe0pRWzpaBl+/VF/BPxapZ6om4Mq4WMMjpowQzLmyIyuEKFbNptM0WRJiSaiOECzOciSMSZL0GA+4kiVGrfCoEpEnf2AXzK6lJa5rael0tPSc33/cDsfq/d13H162qs9fT6SK5HrPue+r7sFX+c6532uK2CMMQIAYIIFXQ8AADA9EUAAACcIIACAEwQQAMAJAggA4AQBBABwggACADhBAAEAnCCAAABOEEAAACdyx+vEW7du1ZNPPqnu7m4tXLhQzz77rBYtWvSlz0skEurq6lJRUZECgcB4DQ8AME6MMerv71dFRYWCwRTzHDMOdu/ebUKhkPnlL39p3n//ffO9733PlJSUmJ6eni99bmdnp5FEo9FotCneOjs7U/7/PmBM5hcjra2t1Y033qif/exnkj6d1VRWVmrdunV6+OGHUz43FouppKREt+i/lKu8TA8NADDOLmhYb+m/1dvbq0gkYn1cxt+CGxoaUnt7u5qampJ9wWBQdXV1amtr+8Lj4/G44vF48s/9/f3/GliecgMEEABMOf+a1nzZxygZL0L4+OOPNTIyorKyslH9ZWVl6u7u/sLjm5ubFYlEkq2ysjLTQwIATELOq+CampoUi8WSrbOz0/WQAAATIONvwV188cXKyclRT0/PqP6enh6Vl5d/4fHhcFjhcDjTwwAATHIZnwGFQiHV1NSopaUl2ZdIJNTS0qJoNJrpywEApqhx+R5QY2OjVq9erRtuuEGLFi3S008/rYGBAd13333jcTkAwBQ0LgG0cuVK/e1vf9Ojjz6q7u5uXXfdddq/f/8XChMAANPXuHwPaCz6+voUiUR0m5ZThg0AU9AFM6yDelmxWEzFxcXWxzmvggMATE8EEADACQIIAOAEAQQAcIIAAgA4QQABAJwggAAAThBAAAAnCCAAgBMEEADACQIIAOAEAQQAcIIAAgA4QQABAJwggAAAThBAAAAnCCAAgBMEEADACQIIAOAEAQQAcIIAAgA4QQABAJwggAAAThBAAAAnCCAAgBMEEADACQIIAOAEAQQAcIIAAgA4QQABAJwggAAAThBAAAAncl0PAJNcIGDpH+fXLiaR4pgZ32sDmBDMgAAAThBAAAAnCCAAgBMEEADACQIIAOAEVXCQgjnWQznFhd4H8kKe3eaTT7z743Hv/pER7/Nne6Wbrbow2+8b+AxmQAAAJwggAIATBBAAwAkCCADgBAEEAHCCKrjpxFJ5FcwP258ze5ZntynwroILnPE+jRkathy4YL92ptgqzlyyVbtRHYdphBkQAMAJAggA4AQBBABwggACADhBAAEAnCCAAABO+A6gN998U3fddZcqKioUCAT00ksvjTpujNGjjz6qOXPmqKCgQHV1dTpx4kSmxotxEMjNtbZEyQzPNlhR5NkCMwo8m0zCu6U14IDPFvTXJoJtrFPtGsAY+P7XNjAwoIULF2rr1q2ex5944glt2bJFzz33nA4fPqwZM2Zo6dKlGhwcHPNgAQDZw/cXUZctW6Zly5Z5HjPG6Omnn9ZPfvITLV++XJL0q1/9SmVlZXrppZf0rW996wvPicfjin9mqf6+vj6/QwIATEEZfb/h1KlT6u7uVl1dXbIvEomotrZWbW1tns9pbm5WJBJJtsrKykwOCQAwSWU0gLq7uyVJZWVlo/rLysqSxz6vqalJsVgs2To7OzM5JADAJOV8LbhwOKxwOMVaZACArJTRACovL5ck9fT0aM6cOcn+np4eXXfddZm8FDIpz/7XYCTf+9hwoffkOVF8kfeJcizbfl+wLEaaTrVWpirYbOdJt2pvPFHVhikso2/BzZ8/X+Xl5WppaUn29fX16fDhw4pGo5m8FABgivM9Azp37pxOnjyZ/POpU6d07NgxlZaWqqqqSuvXr9dPf/pTXX755Zo/f74eeeQRVVRU6O67787kuAEAU5zvADp69Ki+8Y1vJP/c2NgoSVq9erV27typH/3oRxoYGND999+v3t5e3XLLLdq/f7/y8/MzN2oAwJQXMGZy7XTV19enSCSi27RcuYE818PJLpbPC3JKZ1qfMnz1PM/+82XeG9IV/U/Ms990nPLu/8x3wEZx+RmQjcvPgPxuYJfOuYAMuWCGdVAvKxaLqbi42Po41oIDADjhvAwbk8DIiPVQIOH9ankkz/uVtwl7z1oDllfqGX0tbpuh+K1qsz0+1QzL77l8n4dqN2QfZkAAACcIIACAEwQQAMAJAggA4AQBBABwgio4yIzYv98SjNvWavPutq0dl5tn+U6XbaPCVN9VcVQRFgjar2sS4/xaLp316TL53SE/5wf+Q8yAAABOEEAAACcIIACAEwQQAMAJAggA4ARVcJCGh62HggPeq1UHRgq9T1Xo/VcqVDjDsz8x8In3hRP29el8V3dZKsUCtl1a02CrkDOWtfTsFXXeYzIp1uuzD8pyDVtFXaqfOTAOmAEBAJwggAAAThBAAAAnCCAAgBMEEADACargphNL9ViqteAClmMjYe8Kq4Fy779SBbNLvM//8T+8x5SqIstndVeqNdx8SVU1Z612S7FWmx/prAVn4/c51upC1oLD2DADAgA4QQABAJwggAAAThBAAAAnCCAAgBMEEADACcqwIaUqU871Lj0eKvJ+zoWLvE+TyPfekjutAmlbubWtTDpTZdip2K5h26rbUgptW7zUylaeneIa9nNRbo2JxQwIAOAEAQQAcIIAAgA4QQABAJwggAAATlAFh5RbUycKvKvXLnjvsK3hQsuinJZCqnS2ms7UVtrW8yQsFWopFm31f3Hbaz+H22JT7YYJxgwIAOAEAQQAcIIAAgA4QQABAJwggAAATlAFN53Y1vpKweR4v0YZ8S6OUzDufQ3j99JpjNW2Hlvmqubsr9eMrYLM59putu3Dfa8Rlw7WgsMEYwYEAHCCAAIAOEEAAQCcIIAAAE4QQAAAJ6iCQ0q2NdzyPvHuv5Dv3W9yvV/rBG1VXxccVn1NhEztiJqKbb05206pVLthgjEDAgA4QQABAJwggAAAThBAAAAnCCAAgBO+Aqi5uVk33nijioqKdMkll+juu+9WR0fHqMcMDg6qoaFBs2bNUmFhoerr69XT05PRQSNNxni3RMLaAueHPVvugPFug/JsJhjwbIHcXM+WUbb7swkGvdsECAQDns3KJPw368UD3g0YJ77+VbW2tqqhoUGHDh3SgQMHNDw8rDvuuEMDAwPJx2zYsEH79u3Tnj171Nraqq6uLq1YsSLjAwcATG2+Xmru379/1J937typSy65RO3t7fra176mWCymF154Qbt27dLixYslSTt27NBVV12lQ4cO6aabbsrcyAEAU9qY3leIxWKSpNLSUklSe3u7hoeHVVdXl3xMdXW1qqqq1NbW5nmOeDyuvr6+UQ0AkP3SDqBEIqH169fr5ptv1rXXXitJ6u7uVigUUklJyajHlpWVqbu72/M8zc3NikQiyVZZWZnukAAAU0jaAdTQ0KD33ntPu3fvHtMAmpqaFIvFkq2zs3NM5wMATA1plRutXbtWr7zyit58803NnTs32V9eXq6hoSH19vaOmgX19PSovLzc81zhcFjhcDidYSBDrLt5SgoODXv2533i/ZzAiPd5bGvKKc+ytergoHVM1kquhPfrKWN5mRWYhGufTcjOp8Ak4WsGZIzR2rVrtXfvXr3++uuaP3/+qOM1NTXKy8tTS0tLsq+jo0OnT59WNBrNzIgBAFnB1wyooaFBu3bt0ssvv6yioqLk5zqRSEQFBQWKRCJas2aNGhsbVVpaquLiYq1bt07RaJQKOADAKL4CaPv27ZKk2267bVT/jh079J3vfEeStHnzZgWDQdXX1ysej2vp0qXatm1bRgYLAMgevgIo1WcF/5afn6+tW7dq69ataQ8KAJD9WAsOAOAEAQQAcIItuSGlKP0NxIc8+wv+dsGzP15i+Stlefs2nYVHzYh3rXfAsu237/Pk5PgeU6qf4aQzCcvPMT0xAwIAOEEAAQCcIIAAAE4QQAAAJwggAIATVMEh5TbNZjDu2R/6p3d/cNhSYWXpDuRnbiFaa1WbvKvaAjmWxUuHvSv8MmlCFh21Vbv53WabqjmME2ZAAAAnCCAAgBMEEADACQIIAOAEAQQAcIIqOKSuyLJsyZ0TO+/9+KB3hVUgbtmr2yaYYj22hG0NN3/VXWbEXv3nm62SMOD9Gi9g+TnZKvls58koqt0wwZgBAQCcIIAAAE4QQAAAJwggAIATBBAAwAmq4JCSueC9LlpwwLsKLseyvpq1wsq2U2qe/a+msVTU2av5/FXg2SrU0mGtarNf3HIiS5Vdqso1v2u+AROMGRAAwAkCCADgBAEEAHCCAAIAOEEAAQCcoAoOKXdE1bD3WnBmcNCz31p3ZauOswjk2NeCs9Z9WdaIS7munNf5LdV0qarj/D7HWrGX6nfheYE0Kt1Y8w2TBDMgAIATBBAAwAkCCADgBAEEAHCCAAIAOEEVHFKy7RoasOyUamwVVrZKNFvVV4oqOOu5bFVwvivLMve6zPcOp6wFh2mEGRAAwAkCCADgBAEEAHCCAAIAOEEAAQCcIIAAAE5Qho2Upby2MmIzNOT9BJ9bUAdCIV+PT4vt/nyWKdu3/E6D39Jwm1T3wKKjmOSYAQEAnCCAAABOEEAAACcIIACAEwQQAMAJquCQFtsipb4rr9KoUAvkef+1NXF/FXhWtkVNU3G18CeVbpjCmAEBAJwggAAAThBAAAAnCCAAgBMEEADACV8BtH37di1YsEDFxcUqLi5WNBrVq6++mjw+ODiohoYGzZo1S4WFhaqvr1dPT0/GB40JZBKezYyMeDbZWsJ4NjN8wbtdsDerQMC7We/NeLeM/vwydI2JGCswwXwF0Ny5c7Vp0ya1t7fr6NGjWrx4sZYvX673339fkrRhwwbt27dPe/bsUWtrq7q6urRixYpxGTgAYGoLGDO2l1GlpaV68skndc8992j27NnatWuX7rnnHknShx9+qKuuukptbW266aab/qPz9fX1KRKJ6DYtV24gbyxDQybYZhAB79cugaC/x1vZzpOCdYXuiZgp2H5OGVqJm9kOppILZlgH9bJisZiKi4utj0v7M6CRkRHt3r1bAwMDikajam9v1/DwsOrq6pKPqa6uVlVVldra2qznicfj6uvrG9UAANnPdwAdP35chYWFCofDeuCBB7R3715dffXV6u7uVigUUklJyajHl5WVqbu723q+5uZmRSKRZKusrPR9EwCAqcd3AF155ZU6duyYDh8+rAcffFCrV6/WBx98kPYAmpqaFIvFkq2zszPtcwEApg7fa8GFQiFddtllkqSamhodOXJEzzzzjFauXKmhoSH19vaOmgX19PSovLzcer5wOKxwOOx/5JgY1s8evNeCMwnLaxoz7N1v+yzJpPiMxPacnBzvS6eqnMuUTK2BB0wjY/4eUCKRUDweV01NjfLy8tTS0pI81tHRodOnTysajY71MgCALONrBtTU1KRly5apqqpK/f392rVrlw4ePKjXXntNkUhEa9asUWNjo0pLS1VcXKx169YpGo3+xxVwAIDpw1cAnT17Vt/+9rd15swZRSIRLViwQK+99ppuv/12SdLmzZsVDAZVX1+veDyupUuXatu2beMycADA1Dbm7wFlGt8DmiJ8fj9IxrJ/kN/vE6VxjQn5DAhA0rh/DwgAgLFgR1Rklm2m4/Px1mo6SZLPHUv9rlIAYEIwAwIAOEEAAQCcIIAAAE4QQAAAJwggAIATVMEhPZmqILNWqPmsppPs3w+yfm/IZzUdgIxiBgQAcIIAAgA4QQABAJwggAAAThBAAAAnCCAAgBOUYcOtjC4I6rN0m0VKAaeYAQEAnCCAAABOEEAAACcIIACAEwQQAMAJquCQPWzVa7ZqNwBOMQMCADhBAAEAnCCAAABOEEAAACcIIACAE1TBIfv5rY5jjThgQjADAgA4QQABAJwggAAAThBAAAAnCCAAgBNUwWH6Yu04wClmQAAAJwggAIATBBAAwAkCCADgBAEEAHCCKjjg81jzDZgQzIAAAE4QQAAAJwggAIATBBAAwAkCCADgBAEEAHCCAAIAOEEAAQCcIIAAAE4QQAAAJwggAIATYwqgTZs2KRAIaP369cm+wcFBNTQ0aNasWSosLFR9fb16enrGOk4AQJZJO4COHDmi559/XgsWLBjVv2HDBu3bt0979uxRa2ururq6tGLFijEPFACQXdIKoHPnzmnVqlX6+c9/rpkzZyb7Y7GYXnjhBT311FNavHixampqtGPHDv3xj3/UoUOHMjZoAMDUl1YANTQ06M4771RdXd2o/vb2dg0PD4/qr66uVlVVldra2jzPFY/H1dfXN6oBALKf7/2Adu/erbfffltHjhz5wrHu7m6FQiGVlJSM6i8rK1N3d7fn+Zqbm/X444/7HQYAYIrzNQPq7OzUQw89pF//+tfKz8/PyACampoUi8WSrbOzMyPnBQBMbr4CqL29XWfPntX111+v3Nxc5ebmqrW1VVu2bFFubq7Kyso0NDSk3t7eUc/r6elReXm55znD4bCKi4tHNQBA9vP1FtySJUt0/PjxUX333Xefqqur9eMf/1iVlZXKy8tTS0uL6uvrJUkdHR06ffq0otFo5kYNAJjyfAVQUVGRrr322lF9M2bM0KxZs5L9a9asUWNjo0pLS1VcXKx169YpGo3qpptuytyoAQBTnu8ihC+zefNmBYNB1dfXKx6Pa+nSpdq2bVumLwMAmOICxhjjehCf1dfXp0gkotu0XLmBPNfDAQD4dMEM66BeViwWS/m5PmvBAQCcIIAAAE4QQAAAJwggAIATBBAAwAkCCADgBAEEAHCCAAIAOEEAAQCcIIAAAE4QQAAAJwggAIATBBAAwAkCCADgBAEEAHCCAAIAOEEAAQCcIIAAAE4QQAAAJwggAIATua4HgGkiEMjcuYzJ3LkAOMMMCADgBAEEAHCCAAIAOEEAAQCcIIAAAE4QQAAAJyjDRmbZyq0DPl/rmMTYxwJgUmMGBABwggACADhBAAEAnCCAAABOEEAAACeogkN6fFa7BYLejzcJFhYFpitmQAAAJwggAIATBBAAwAkCCADgBAEEAHCCKjikZql2C+TkeD/e1m87vUY8+413t1vpbCvO9uGAFTMgAIATBBAAwAkCCADgBAEEAHCCAAIAOEEVHFJWd9mq3QKhkHd/rvdfKXPhgnf/iKXcbSKqx9KpanN1DarpkIWYAQEAnCCAAABOEEAAACcIIACAEwQQAMAJXwH02GOPKRAIjGrV1dXJ44ODg2poaNCsWbNUWFio+vp69fT0ZHzQyKxATo69FRR4txkzPJsK8j1bIDfXsykQ9G4ZvcGApVmunUaz/fwA2Pn+l37NNdfozJkzyfbWW28lj23YsEH79u3Tnj171Nraqq6uLq1YsSKjAwYAZAff3wPKzc1VeXn5F/pjsZheeOEF7dq1S4sXL5Yk7dixQ1dddZUOHTqkm266aeyjBQBkDd8zoBMnTqiiokKXXnqpVq1apdOnT0uS2tvbNTw8rLq6uuRjq6urVVVVpba2Nuv54vG4+vr6RjUAQPbzFUC1tbXauXOn9u/fr+3bt+vUqVO69dZb1d/fr+7uboVCIZWUlIx6TllZmbq7u63nbG5uViQSSbbKysq0bgQAMLX4egtu2bJlyf9esGCBamtrNW/ePL344osqKChIawBNTU1qbGxM/rmvr48QAoBpYExrwZWUlOiKK67QyZMndfvtt2toaEi9vb2jZkE9PT2enxn9WzgcVjgcHssw8J+y7W5qWddNkoKRYs9+U3SR9xMSljXLbP2ffGK9tm9+d2+1nidzVXgBn5c2tp+TEmMey/9ehHXlMDmM6V/auXPn9NFHH2nOnDmqqalRXl6eWlpaksc7Ojp0+vRpRaPRMQ8UAJBdfM2AfvjDH+quu+7SvHnz1NXVpY0bNyonJ0f33nuvIpGI1qxZo8bGRpWWlqq4uFjr1q1TNBqlAg4A8AW+Auivf/2r7r33Xv3973/X7Nmzdcstt+jQoUOaPXu2JGnz5s0KBoOqr69XPB7X0qVLtW3btnEZOABgagsYM7neEO7r61MkEtFtWq7cQJ7r4WQXy2ckwRQFJMHSmZ79vj8D+kfM++H//Kf3+S37B6U0CT8DkvH32Y31MyCf50l9kUn1Tx5Z6IIZ1kG9rFgspuJi78+RJdaCAwA4QgABAJxgS+5pxLq9dkG+9Tkj5d5vwZ2f4/0WXKhv2Ls/PuR9ActbcOmw3p9lm3AFfb7+SrW9tmVrcWMsz7G81WYt27acx162Lfvbdrb74K05TDBmQAAAJwggAIATBBAAwAkCCADgBAEEAHCCKrjpxPIFy0BxkfUpA1UzPPt7/493uVbuee+/UmXnLNc4Yyn7SvFFVGtVm60KLt+y2K3tC6rpfBk0zzKmEctzbFVzln4lLL+7oH1MxvIcO8u5qI7DOGEGBABwggACADhBAAEAnCCAAABOEEAAACeogstG1m0JvF9vJIrs2zH8/SrvSrHwon949sf+EvHsj/w/77XjLrL0j8Tj1jHZBG1r2lm2fA/Y1oILWbYBuWCpUJNkhrzXwLPuKGKr8huyrJnnd605SYGAd/WatdLOfiJ/j6dqDv8hZkAAACcIIACAEwQQAMAJAggA4AQBBABwgiq4bGSrQrJUfSUuCllPNXT1ec/+/f/3F579D198l2d/x3vVnv0zwvZr+2Us664FLWvHmYvsO8F6slXHSQrY1oKz7ARrqxMLBC0VZ8OWqrlUFWfD3pV5dt4Vj9aqOb87rqZC5dy0xAwIAOAEAQQAcIIAAgA4QQABAJwggAAATlAFN40YSyVVcNC+++jIoHeVWlVuoWf/d8re8ux/aM5Vnv2Ji2d69gf+2Wsdk7GsExewVH1Z12lLUdXmeR5bpZskY6nms9WDBYKWn/mI92tCY1vfL8XOsSZh2SHWUg1pq3YLWNebs1TN2XaU/fSgdz/rzU1LzIAAAE4QQAAAJwggAIATBBAAwAkCCADgBFVw04gZtuy2+d7/WJ9z9U+/4tlfe/BBz/7+Ku9qplzLBqcDlxV79hcOV1rHZLp6PPsT/f3eTxgc9O7/u3cVV9BSHRfI995Z9dODlsoyyzp0SljWV7NVkFmqx2zr30myrv2nRIrneLHcm/3h9vObhOVctuo460WomssGzIAAAE4QQAAAJwggAIATBBAAwAkCCADgBAEEAHCCMmzIpFjQ8sKpv3j2l9j6/V7cUk47kqL017ZtdcBW8mxhWzQzYVns1FrODSAtzIAAAE4QQAAAJwggAIATBBAAwAkCCADgBFVwcMu2SKSxLNYp/+tWApicmAEBAJwggAAAThBAAAAnCCAAgBOTrgjB/OtD6QsaltjEEACmnAsalvS//z+3mXQB1P+vbZXf0n87HgkAYCz6+/sViUSsxwPmyyJqgiUSCXV1damoqEj9/f2qrKxUZ2eniouLXQ9twvT19XHf0+S+p+M9S9PzvqfTPRtj1N/fr4qKCgWD9k96Jt0MKBgMau7cuZKkwL9WSi4uLs76X5gX7nv6mI73LE3P+54u95xq5vNvFCEAAJwggAAATkzqAAqHw9q4caPC4bDroUwo7nv63Pd0vGdpet73dLznLzPpihAAANPDpJ4BAQCyFwEEAHCCAAIAOEEAAQCcIIAAAE5M6gDaunWrvvrVryo/P1+1tbX605/+5HpIGfXmm2/qrrvuUkVFhQKBgF566aVRx40xevTRRzVnzhwVFBSorq5OJ06ccDPYDGlubtaNN96ooqIiXXLJJbr77rvV0dEx6jGDg4NqaGjQrFmzVFhYqPr6evX09DgacWZs375dCxYsSH4LPhqN6tVXX00ez8Z7/rxNmzYpEAho/fr1yb5svO/HHntMgUBgVKuurk4ez8Z7TtekDaDf/va3amxs1MaNG/X2229r4cKFWrp0qc6ePet6aBkzMDCghQsXauvWrZ7Hn3jiCW3ZskXPPfecDh8+rBkzZmjp0qUaHByc4JFmTmtrqxoaGnTo0CEdOHBAw8PDuuOOOzQwMJB8zIYNG7Rv3z7t2bNHra2t6urq0ooVKxyOeuzmzp2rTZs2qb29XUePHtXixYu1fPlyvf/++5Ky854/68iRI3r++ee1YMGCUf3Zet/XXHONzpw5k2xvvfVW8li23nNazCS1aNEi09DQkPzzyMiIqaioMM3NzQ5HNX4kmb179yb/nEgkTHl5uXnyySeTfb29vSYcDpvf/OY3DkY4Ps6ePWskmdbWVmPMp/eYl5dn9uzZk3zMn//8ZyPJtLW1uRrmuJg5c6b5xS9+kfX33N/fby6//HJz4MAB8/Wvf9089NBDxpjs/V1v3LjRLFy40PNYtt5zuiblDGhoaEjt7e2qq6tL9gWDQdXV1amtrc3hyCbOqVOn1N3dPepnEIlEVFtbm1U/g1gsJkkqLS2VJLW3t2t4eHjUfVdXV6uqqipr7ntkZES7d+/WwMCAotFo1t9zQ0OD7rzzzlH3J2X37/rEiROqqKjQpZdeqlWrVun06dOSsvue0zHpVsOWpI8//lgjIyMqKysb1V9WVqYPP/zQ0agmVnd3tyR5/gz+fWyqSyQSWr9+vW6++WZde+21kj6971AopJKSklGPzYb7Pn78uKLRqAYHB1VYWKi9e/fq6quv1rFjx7L2nnfv3q23335bR44c+cKxbP1d19bWaufOnbryyit15swZPf7447r11lv13nvvZe09p2tSBhCmh4aGBr333nuj3h/PZldeeaWOHTumWCym3/3ud1q9erVaW1tdD2vcdHZ26qGHHtKBAweUn5/vejgTZtmyZcn/XrBggWprazVv3jy9+OKLKigocDiyyWdSvgV38cUXKycn5wuVIT09PSovL3c0qon17/vM1p/B2rVr9corr+iNN95I7v8kfXrfQ0ND6u3tHfX4bLjvUCikyy67TDU1NWpubtbChQv1zDPPZO09t7e36+zZs7r++uuVm5ur3Nxctba2asuWLcrNzVVZWVlW3vfnlZSU6IorrtDJkyez9nedrkkZQKFQSDU1NWppaUn2JRIJtbS0KBqNOhzZxJk/f77Ky8tH/Qz6+vp0+PDhKf0zMMZo7dq12rt3r15//XXNnz9/1PGamhrl5eWNuu+Ojg6dPn16St+3l0QioXg8nrX3vGTJEh0/flzHjh1LthtuuEGrVq1K/nc23vfnnTt3Th999JHmzJmTtb/rtLmugrDZvXu3CYfDZufOneaDDz4w999/vykpKTHd3d2uh5Yx/f395p133jHvvPOOkWSeeuop884775i//OUvxhhjNm3aZEpKSszLL79s3n33XbN8+XIzf/58c/78eccjT9+DDz5oIpGIOXjwoDlz5kyyffLJJ8nHPPDAA6aqqsq8/vrr5ujRoyYajZpoNOpw1GP38MMPm9bWVnPq1Cnz7rvvmocfftgEAgHz+9//3hiTnffs5bNVcMZk533/4Ac/MAcPHjSnTp0yf/jDH0xdXZ25+OKLzdmzZ40x2XnP6Zq0AWSMMc8++6ypqqoyoVDILFq0yBw6dMj1kDLqjTfeMJK+0FavXm2M+bQU+5FHHjFlZWUmHA6bJUuWmI6ODreDHiOv+5VkduzYkXzM+fPnzfe//30zc+ZMc9FFF5lvfvOb5syZM+4GnQHf/e53zbx580woFDKzZ882S5YsSYaPMdl5z14+H0DZeN8rV640c+bMMaFQyHzlK18xK1euNCdPnkwez8Z7Thf7AQEAnJiUnwEBALIfAQQAcIIAAgA4QQABAJwggAAAThBAAAAnCCAAgBMEEADACQIIAOAEAQQAcIIAAgA48f8BQyiUbDWvQCEAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "edges_x1 = torch.from_numpy(inhist.axes[0].edges)\n",
    "edges_x2 = torch.from_numpy(inhist.axes[1].edges)\n",
    "\n",
    "centers_x1 = torch.diff(edges_x1) / 2 + edges_x1[:-1]\n",
    "centers_x2 = torch.diff(edges_x2) / 2 + edges_x2[:-1]\n",
    "centers_x1.shape\n",
    "\n",
    "bin_values = torch.from_numpy(inhist.values()).T\n",
    "dx,dy =torch.gradient(bin_values,spacing=(centers_x1,centers_x2))\n",
    "D = dx**2 + dy**2\n",
    "print(g.shape)\n",
    "fig,ax=plt.subplots()\n",
    "ax.imshow(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "2a4d6e8b-c54e-4710-adad-966ad6949204",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LinearTransform(tensor([2.6500e+03, 9.0000e-01]), tensor([0., 0.]))\n",
      "LinearTransform(tensor([5342.7125]), tensor([4047.3054]))\n"
     ]
    }
   ],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "def bumpCut(X, Y):\n",
    "    m = Y > (1 - 200 / X)\n",
    "    return m\n",
    "\n",
    "window_func = windowing.EllipseWindow([1300,0.5],[160,0.06])\n",
    "#window_func = windowing.EllipseWindow([0.5,0.5],[0.1,0.1])\n",
    "\n",
    "min_counts = 50\n",
    "train_data,*_ = regression.makeRegressionData(\n",
    "    inhist, window_func, domain_mask_function=bumpCut, exclude_less=min_counts\n",
    ")\n",
    "test_data, domain_mask,reshape,*_ = regression.makeRegressionData(\n",
    "    inhist,\n",
    "    None,\n",
    "    get_mask=True,\n",
    "    domain_mask_function=bumpCut,\n",
    "    exclude_less=min_counts,\n",
    "    get_reshape_function=True,\n",
    ")\n",
    "train_transform = transformations.getNormalizationTransform(train_data)\n",
    "normalized_train_data = train_transform.transform(train_data)\n",
    "normalized_test_data = train_transform.transform(test_data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "3cb96406-4d5a-4291-849a-08f553af4755",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "unflatten: Provided sizes [60, 60] don't multiply up to the size of dim 1 (2) in the input tensor",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[30], line 5\u001b[0m\n\u001b[1;32m      3\u001b[0m X \u001b[38;5;241m=\u001b[39m normalized_test_data\u001b[38;5;241m.\u001b[39mX\n\u001b[1;32m      4\u001b[0m \u001b[38;5;66;03m#torch.gradient(Y,X)\u001b[39;00m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[43mreshape\u001b[49m\u001b[43m(\u001b[49m\u001b[43mX\u001b[49m\u001b[43m)\u001b[49m)\n",
      "File \u001b[0;32m/srv/fitting/regression.py:84\u001b[0m, in \u001b[0;36mmakeRegressionData.<locals>.<lambda>\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m     82\u001b[0m     ret \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m*\u001b[39mret, centers_mask)\n\u001b[1;32m     83\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m get_reshape_function:\n\u001b[0;32m---> 84\u001b[0m     ret \u001b[38;5;241m=\u001b[39m (\u001b[38;5;241m*\u001b[39mret, \u001b[38;5;28;01mlambda\u001b[39;00m x: \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43munflatten\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbin_values\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mshape\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m     85\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m ret\n",
      "\u001b[0;31mRuntimeError\u001b[0m: unflatten: Provided sizes [60, 60] don't multiply up to the size of dim 1 (2) in the input tensor"
     ]
    }
   ],
   "source": [
    "normalized_test_data.X\n",
    "Y = normalized_test_data.Y\n",
    "X = normalized_test_data.X\n",
    "#torch.gradient(Y,X)\n",
    "print(reshape(X))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c4b2a8a-b439-4997-9354-61e3b9225b9a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'tuple' object has no attribute 'X'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[10], line 20\u001b[0m\n\u001b[1;32m     10\u001b[0m train_data \u001b[38;5;241m=\u001b[39m regression\u001b[38;5;241m.\u001b[39mmakeRegressionData(\n\u001b[1;32m     11\u001b[0m     inhist, window_func, domain_mask_function\u001b[38;5;241m=\u001b[39mbumpCut, exclude_less\u001b[38;5;241m=\u001b[39mmin_counts\n\u001b[1;32m     12\u001b[0m )\n\u001b[1;32m     13\u001b[0m test_data, domain_mask \u001b[38;5;241m=\u001b[39m regression\u001b[38;5;241m.\u001b[39mmakeRegressionData(\n\u001b[1;32m     14\u001b[0m     inhist,\n\u001b[1;32m     15\u001b[0m     \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     18\u001b[0m     exclude_less\u001b[38;5;241m=\u001b[39mmin_counts,\n\u001b[1;32m     19\u001b[0m )\n\u001b[0;32m---> 20\u001b[0m train_transform \u001b[38;5;241m=\u001b[39m \u001b[43mtransformations\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mgetNormalizationTransform\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtrain_data\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     21\u001b[0m normalized_train_data \u001b[38;5;241m=\u001b[39m train_transform\u001b[38;5;241m.\u001b[39mtransform(train_data)\n\u001b[1;32m     22\u001b[0m normalized_test_data \u001b[38;5;241m=\u001b[39m train_transform\u001b[38;5;241m.\u001b[39mtransform(test_data)\n",
      "File \u001b[0;32m/srv/fitting/transformations.py:98\u001b[0m, in \u001b[0;36mgetNormalizationTransform\u001b[0;34m(dv, scale)\u001b[0m\n\u001b[1;32m     97\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mgetNormalizationTransform\u001b[39m(dv, scale\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1.0\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DataTransformation:\n\u001b[0;32m---> 98\u001b[0m     X, Y, V, E \u001b[38;5;241m=\u001b[39m \u001b[43mdv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mX\u001b[49m, dv\u001b[38;5;241m.\u001b[39mY, dv\u001b[38;5;241m.\u001b[39mV, dv\u001b[38;5;241m.\u001b[39mE\n\u001b[1;32m    100\u001b[0m     max_x, min_x \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mmax(X, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\u001b[38;5;241m.\u001b[39mvalues, torch\u001b[38;5;241m.\u001b[39mmin(X, axis\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0\u001b[39m)\u001b[38;5;241m.\u001b[39mvalues\n\u001b[1;32m    101\u001b[0m     max_y, min_y \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mmax(Y), torch\u001b[38;5;241m.\u001b[39mmin(Y)\n",
      "\u001b[0;31mAttributeError\u001b[0m: 'tuple' object has no attribute 'X'"
     ]
    }
   ],
   "source": [
    "kernel = SK(gpytorch.kernels.RBFKernel(ard_num_dims=2)) \n",
    "\n",
    "# kernel = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4,ard_num_dims=2)\n",
    "#kernel = SK(gpytorch.kernels.MaternKernel(ard_num_dims=2)) * SK(gpytorch.kernels.MaternKernel(nu=1.5, ard_num_dims=2))* SK(gpytorch.kernels.MaternKernel(nu=0.5, ard_num_dims=2))\n",
    "use_cuda = True\n",
    "\n",
    "\n",
    "if torch.cuda.is_available() and use_cuda:\n",
    "    train = normalized_train_data.toGpu()\n",
    "    norm_test = normalized_test_data.toGpu()\n",
    "    print(\"USING CUDA\")\n",
    "else:\n",
    "    train = normalized_train_data\n",
    "    norm_test = normalized_test_data\n",
    "    \n",
    "\n",
    "#trainX, trainY, trainV = batchify(train.X,train.Y,train.V, size=800, batches=10)\n",
    "\n",
    "trainX, trainY, trainV = train.X, train.Y, train.V\n",
    "\n",
    "#kernel = SK(gpytorch.kernels.MaternKernel(ard_num_dims=2)) * SK(gpytorch.kernels.RBFKernel(ard_num_dims=2))\n",
    "likelihood = gpytorch.likelihoods.FixedNoiseGaussianLikelihood(\n",
    "    noise=trainV*4,\n",
    "    learn_additional_noise=False,\n",
    ")\n",
    "\n",
    "#model = GPRegressionModel(trainX, trainY, likelihood)\n",
    "#model = NonStatGP(normalized_train_data.X, normalized_train_data.Y, likelihood,)\n",
    "model = models.ExactAnyKernelModel(trainX, trainY, likelihood, kernel=kernel)\n",
    "print(model)\n",
    "lr = 0.05\n",
    "if torch.cuda.is_available() and use_cuda:\n",
    "    model = model.cuda()\n",
    "    likelihood = likelihood.cuda()\n",
    "\n",
    "# if hasattr(model.covar_module, \"initialize_from_data\"):\n",
    "#     model.covar_module.initialize_from_data(train.X, train.Y)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.1)  # Includes GaussianLikelihood parameters\n",
    "\n",
    "# \"Loss\" for GPs - the marginal log likelihood\n",
    "mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)\n",
    "\n",
    "for i in range(400):\n",
    "    # Zero gradients from previous iteration\n",
    "    optimizer.zero_grad()\n",
    "    # Output from model\n",
    "    output = model(trainX)\n",
    "    # Calc loss and backprop gradients\n",
    "    l = -mll(output, trainY)\n",
    "    mse = gpytorch.metrics.mean_squared_error(output, trainY,  squared=True)\n",
    "    #print(l)\n",
    "    loss = l\n",
    "    loss.backward()\n",
    "    if i % 20 == 0 :     \n",
    "        print('Iter %d/%d - Loss: %.3f' % (i + 1, 200, loss.item()))\n",
    "    \n",
    "    optimizer.step()\n",
    "\n",
    "# model, likelihood, evidence = regression.optimizeHyperparams(\n",
    "#     model,\n",
    "#     likelihood,\n",
    "#     train,\n",
    "#     bar=False,\n",
    "#     iterations=800,\n",
    "#     lr=lr,\n",
    "#     get_evidence=True,\n",
    "# )\n",
    "if torch.cuda.is_available() and use_cuda:\n",
    "    model = model.cpu()\n",
    "    likelihood = likelihood.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "693f3576-d0f3-4388-b63b-4a02c28f6467",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.covar_module.base_kernel.lengthscale = torch.tensor([[0.1,0.1]])\n",
    "model.covar_module.outputscale = torch.tensor(0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "033b4928-0001-4b0e-bec6-6d52b783db17",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe6dc2c5-4c5c-4549-8d01-460641866ace",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for param_name, param in model.named_parameters():\n",
    "    print(f'Parameter name: {param_name:42} value = {param.round(decimals=2).tolist()}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54e9104c-8064-468e-9a44-4c35357b1a65",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_dist = regression.getPrediction(\n",
    "    model, likelihood, normalized_test_data\n",
    ")\n",
    "print(pred_dist)\n",
    "\n",
    "\n",
    "with gpytorch.settings.cholesky_max_tries(30):\n",
    "    psd_pred_dist = fixMVN(pred_dist)\n",
    "pred_dist = type(pred_dist)(\n",
    "    psd_pred_dist.mean, psd_pred_dist.covariance_matrix.to_dense()\n",
    ")\n",
    "slope = train_transform.transform_y.slope\n",
    "intercept = train_transform.transform_y.intercept\n",
    "pred_dist = affineTransformMVN(psd_pred_dist, slope, intercept)\n",
    "\n",
    "\n",
    "to_use = test_data\n",
    "to_use = regression.DataValues(\n",
    "    to_use.X,#.unsqueeze(0).repeat(4,1,1),\n",
    "    to_use.Y,#.unsqueeze(0).repeat(4,1),\n",
    "    to_use.V,#.unsqueeze(0).repeat(4,1),\n",
    "    to_use.E,\n",
    ")\n",
    "\n",
    "\n",
    "\n",
    "pred_data = regression.DataValues(\n",
    "    to_use.X,\n",
    "    pred_dist.mean,\n",
    "    pred_dist.variance,\n",
    "    to_use.E,\n",
    ")\n",
    "\n",
    "mask = regression.getBlindedMask(pred_data.X, pred_data.Y, to_use.Y, to_use.V, window_func)\n",
    "print(mask)\n",
    "#gbm = (slice(None),torch.ones_like( test_data.Y > 50))\n",
    "gbm = torch.ones_like(test_data.Y > 50)\n",
    "\n",
    "x =     torch.sum((pred_data.Y[gbm] - to_use.Y[gbm]) ** 2/ to_use.V[gbm], dim=0)\n",
    "global_chi2_bins = torch.sum((pred_data.Y[gbm] - to_use.Y[gbm]) ** 2/ to_use.V[gbm], dim=0)/ to_use.Y[gbm].shape[0]\n",
    "\n",
    "print(f\"Global Chi2/bins = {global_chi2_bins}\")\n",
    "ok = True\n",
    "\n",
    "\n",
    "mask = regression.getBlindedMask(\n",
    "    pred_data.X, pred_data.Y, test_data.Y, test_data.V, window_func\n",
    ")\n",
    "wm= (slice(None), mask)\n",
    "print(mask.shape)\n",
    "bpred_mean = pred_data.Y[mask]\n",
    "obs_mean = to_use.Y[mask]\n",
    "obs_var = to_use.V[mask]\n",
    "print(bpred_mean)\n",
    "chi2 = torch.sum((obs_mean - bpred_mean) ** 2 / obs_var) / torch.count_nonzero(\n",
    "    mask\n",
    ")\n",
    "avg_pull = torch.sum(\n",
    "    torch.abs((obs_mean - bpred_mean)) / torch.sqrt(obs_var),dim=0\n",
    ") / torch.count_nonzero(mask)\n",
    "\n",
    "print(f\"Avg Abs pull = {avg_pull}\")\n",
    "print(f\"Chi^2/bins = {chi2}\")\n",
    "\n",
    "print(f\"Global Chi2/bins = {global_chi2_bins}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1329a8ba-077a-4872-bee0-696516289d3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "if hasattr(model.covar_module, \"inducing_points\"):\n",
    "    i = model.covar_module.inducing_points.detach().numpy()\n",
    "else:\n",
    "    i = None\n",
    "plots = makeDiagnosticPlots(pred_data,\n",
    "                            test_data, \n",
    "                            train_data,\n",
    "                            mask,\n",
    "                            inducing_points=i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbabad09-371d-4887-a53c-7856ba989195",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "fe = model.feature_extractor\n",
    "#for param_name, param in model.named_parameters():\n",
    "#    print(f'Parameter name: {param_name:42} value = {param.round(decimals=2).tolist()}')\n",
    "\n",
    "s = model.scale_to_bounds\n",
    "def t(x):\n",
    "    return fe(x)\n",
    "\n",
    "T = t(test_data.X).detach()\n",
    "print(T[:,0] / T[:,1])\n",
    "fig,ax=plt.subplots()\n",
    "ax.scatter(T[:,0], T[:,1], c=test_data.Y, cmap='hsv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7276b243-73e4-408b-8cb2-c120589a0f4a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55f5b78d-e873-4e48-9ccf-20040c7ad9a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "   logger.info(\"Done with loop\")\n",
    "            print(\"Done with loop\")\n",
    "            raw_pred_dist = regression.getPrediction(\n",
    "                model, likelihood, normalized_test_data\n",
    "            )\n",
    "            with gpytorch.settings.cholesky_max_tries(30):\n",
    "                psd_pred_dist = fit_utils.fixMVN(raw_pred_dist)\n",
    "            raw_pred_dist = type(raw_pred_dist)(\n",
    "                psd_pred_dist.mean, psd_pred_dist.covariance_matrix.to_dense()\n",
    "            )\n",
    "            slope = train_transform.transform_y.slope\n",
    "            intercept = train_transform.transform_y.intercept\n",
    "            pred_dist = fit_utils.affineTransformMVN(psd_pred_dist, slope, intercept)\n",
    "\n",
    "            pred_data = regression.DataValues(\n",
    "                test_data.X,\n",
    "                pred_dist.mean,\n",
    "                pred_dist.variance,\n",
    "                test_data.E,\n",
    "            )\n",
    "\n",
    "            good_bin_mask = test_data.Y > 500\n",
    "            global_chi2_bins = float(\n",
    "                torch.sum(\n",
    "                    (pred_data.Y[good_bin_mask] - test_data.Y[good_bin_mask]) ** 2\n",
    "                    / test_data.V[good_bin_mask]\n",
    "                )\n",
    "                / test_data.Y[good_bin_mask].shape[0]\n",
    "            )\n",
    "            print(f\"Global Chi2/bins = {global_chi2_bins}\")\n",
    "            if global_chi2_bins < 1.5:\n",
    "                ok = True\n",
    "            else:\n",
    "                logger.warning(\"Bad global Chi2, retrying\")\n",
    "            ok = True\n",
    "\n",
    "        except (\n",
    "            linear_operator.utils.errors.NanError,\n",
    "            linear_operator.utils.errors.NotPSDError,\n",
    "        ) as e:\n",
    "            lr = lr + random.random() / 1000\n",
    "            logger.warning(f\"CHOLESKY FAILED: retrying with lr={round(lr,3)}\")\n",
    "            logger.warning(e)\n",
    "\n",
    "    logger.warning(\"Done training\")\n",
    "\n",
    "    data = {\n",
    "        \"evidence\": evidence,\n",
    "        \"global_chi2/bins\": global_chi2_bins,\n",
    "        \"model_string\": str(model),\n",
    "    }\n",
    "\n",
    "    if window_func:\n",
    "        mask = regression.getBlindedMask(\n",
    "            pred_data.X, pred_data.Y, test_data.Y, test_data.V, window_func\n",
    "        )\n",
    "        bpred_mean = pred_data.Y[mask]\n",
    "        obs_mean = test_data.Y[mask]\n",
    "        obs_var = test_data.V[mask]\n",
    "        chi2 = torch.sum((obs_mean - bpred_mean) ** 2 / obs_var) / torch.count_nonzero(\n",
    "            mask\n",
    "        )\n",
    "        avg_pull = torch.sum(\n",
    "            torch.abs((obs_mean - bpred_mean)) / torch.sqrt(obs_var)\n",
    "        ) / torch.count_nonzero(mask)\n",
    "\n",
    "        data.update(\n",
    "            {\n",
    "                \"chi2_blinded\": float(chi2),\n",
    "                \"avg_abs_pull\": float(avg_pull),\n",
    "            }\n",
    "        )\n",
    "        print(f\"Avg Abs pull = {avg_pull}\")\n",
    "        print(f\"Chi^2/bins = {chi2}\")\n",
    "    else:\n",
    "        mask = None\n",
    "    print(f\"Global Chi2/bins = {global_chi2_bins}\")\n",
    "\n",
    "    save_dir = Path(save_dir)\n",
    "    save_dir.mkdir(exist_ok=True, parents=True)\n",
    "    if True:\n",
    "        diagnostic_plots = makeDiagnosticPlots(pred_data, test_data, train_data, mask)\n",
    "        saveDiagnosticPlots(diagnostic_plots, save_dir)\n",
    "        # makeSlicePlots(pred_data, test_data, inhist, window_func, 0, save_dir)\n",
    "        # makeSlicePlots(pred_data, test_data, inhist, window_func, 1, save_dir)\n",
    "\n",
    "    model_dict = model.state_dict()\n",
    "\n",
    "    save_data = RegressionModel(\n",
    "        input_data=inhist,\n",
    "        window=window_func,\n",
    "        domain_mask=domain_mask,\n",
    "        train_data=train_data,\n",
    "        test_data=test_data,\n",
    "        trained_model=model,\n",
    "        raw_posterior_dist=raw_pred_dist,\n",
    "        posterior_dist=pred_dist,\n",
    "    )\n",
    "    torch.save(save_data, save_dir / \"train_model.pth\")\n",
    "    torch.save(save_data, save_dir / \"model_dict.pth\")\n",
    "    torch.save(pred_dist, save_dir / \"posterior_latent.pth\")\n",
    "    with open(save_dir / \"info.json\", \"w\") as f:\n",
    "        json.dump(data, f)\n",
    "    return save_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "575ca4a7-4145-4645-bbad-728bd596a163",
   "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}