CAMELS_CKA_HI_After_ICML_Train_Classifier (1).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# 1. Unchanged code"
      ],
      "metadata": {
        "id": "mnULovMXMDsq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Unchanged code\n",
        "!pip install optuna==2.4.0\n",
        "\n",
        "import tqdm\n",
        "import numpy as np\n",
        "\n",
        "import numpy as np\n",
        "import optuna\n",
        "from google.colab import drive\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import sys,os\n",
        "import random\n",
        "from torch.utils.data import DataLoader\n",
        "\n",
        "import shutil\n",
        "import time\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
        "\n",
        "import matplotlib as mpl\n",
        "mpl.rcParams['figure.dpi'] = 200\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "    print(\"CUDA Available\")\n",
        "    device = torch.device('cuda')\n",
        "else:\n",
        "    print('CUDA Not Available')\n",
        "    device = torch.device('cpu')\n",
        "\n",
        "\n",
        "def unbiased_HSIC(K, L):\n",
        "  '''Computes an unbiased estimator of HISC. This is equation (2) from the paper'''\n",
        "\n",
        "  #create the unit **vector** filled with ones\n",
        "  n = K.shape[0]\n",
        "  ones = np.ones(shape=(n))\n",
        "\n",
        "  #fill the diagonal entries with zeros\n",
        "  np.fill_diagonal(K, val=0) #this is now K_tilde\n",
        "  np.fill_diagonal(L, val=0) #this is now L_tilde\n",
        "\n",
        "  #first part in the square brackets\n",
        "  trace = np.trace(np.dot(K, L))\n",
        "\n",
        "  #middle part in the square brackets\n",
        "  nominator1 = np.dot(np.dot(ones.T, K), ones)\n",
        "  nominator2 = np.dot(np.dot(ones.T, L), ones)\n",
        "  denominator = (n-1)*(n-2)\n",
        "  middle = np.dot(nominator1, nominator2) / denominator\n",
        "\n",
        "\n",
        "  #third part in the square brackets\n",
        "  multiplier1 = 2/(n-2)\n",
        "  multiplier2 = np.dot(np.dot(ones.T, K), np.dot(L, ones))\n",
        "  last = multiplier1 * multiplier2\n",
        "\n",
        "  #complete equation\n",
        "  unbiased_hsic = 1/(n*(n-3)) * (trace + middle - last)\n",
        "\n",
        "  return unbiased_hsic\n",
        "\n",
        "def CKA(X, Y):\n",
        "  '''Computes the CKA of two matrices. This is equation (1) from the paper'''\n",
        "\n",
        "  nominator = unbiased_HSIC(np.dot(X, X.T), np.dot(Y, Y.T))\n",
        "  denominator1 = unbiased_HSIC(np.dot(X, X.T), np.dot(X, X.T))\n",
        "  denominator2 = unbiased_HSIC(np.dot(Y, Y.T), np.dot(Y, Y.T))\n",
        "\n",
        "  cka = nominator/np.sqrt(denominator1*denominator2)\n",
        "\n",
        "  return cka\n",
        "\n",
        "def calculate_CKA_for_two_matrices(activationA, activationB):\n",
        "  '''Takes two activations A and B and computes the linear CKA to measure their similarity'''\n",
        "\n",
        "  #unfold the activations, that is make a (n, h*w*c) representation\n",
        "  shape = activationA.shape\n",
        "  activationA = np.reshape(activationA, newshape=(shape[0], np.prod(shape[1:])))\n",
        "\n",
        "  shape = activationB.shape\n",
        "  activationB = np.reshape(activationB, newshape=(shape[0], np.prod(shape[1:])))\n",
        "\n",
        "  #calculate the CKA score\n",
        "  cka_score = CKA(activationA, activationB)\n",
        "\n",
        "  del activationA\n",
        "  del activationB\n",
        "\n",
        "  return cka_score\n",
        "\n",
        "\n",
        "# def get_all_layer_outputs_fn(model):\n",
        "#   '''Builds and returns function that returns the output of every (intermediate) layer'''\n",
        "\n",
        "#   return tf.keras.backend.function([model.layers[0].input],\n",
        "#                                   [l.output for l in model.layers[1:]])\n",
        "\n",
        "def compare_activations(modelA, modelB, data_batch):\n",
        "  '''\n",
        "  Calculate a pairwise comparison of hidden representations and return a matrix\n",
        "  '''\n",
        "\n",
        "  _, intermediate_outputs_A = modelA(data_batch)\n",
        "  _, intermediate_outputs_B = modelB(data_batch)\n",
        "\n",
        "  #create a placeholder array\n",
        "  result_array = np.zeros(shape=(len(intermediate_outputs_A), len(intermediate_outputs_B)))\n",
        "\n",
        "  i = 0\n",
        "  for outputA in tqdm.tqdm_notebook(intermediate_outputs_A):\n",
        "    j = 0\n",
        "    for outputB in tqdm.tqdm_notebook(intermediate_outputs_B):\n",
        "      print(outputA.shape, outputB.shape)\n",
        "      cka_score = calculate_CKA_for_two_matrices(outputA, outputB)\n",
        "      result_array[i, j] = cka_score\n",
        "      j+=1\n",
        "    i+= 1\n",
        "\n",
        "  return result_array\n",
        "\n",
        "# This routine returns the data loader need to train the network\n",
        "def create_dataset_multifield(mode, seed, fmaps, fparams, batch_size, splits, fmaps_norm,\n",
        "                              rot_flip_in_mem=True, shuffle=True, verbose=False):\n",
        "\n",
        "    # whether rotations and flippings are kept in memory\n",
        "    if rot_flip_in_mem:\n",
        "        data_set = make_dataset_multifield(mode, seed, fmaps, fparams, splits, fmaps_norm, verbose)\n",
        "    else:\n",
        "        data_set = make_dataset_multifield2(mode, seed, fmaps, fparams, splits, fmaps_norm, verbose)\n",
        "\n",
        "    def seed_worker(worker_id):\n",
        "        worker_seed = torch.initial_seed() % 2**32\n",
        "        np.random.seed(worker_seed)\n",
        "        random.seed(worker_seed)\n",
        "\n",
        "    g = torch.Generator()\n",
        "    g.manual_seed(0)\n",
        "\n",
        "    data_loader = DataLoader(dataset=data_set, batch_size=batch_size, shuffle=shuffle, worker_init_fn=seed_worker, generator=g)\n",
        "    return data_loader\n",
        "\n",
        "\n",
        "# This class creates the dataset. Rotations and flippings are precompued and stored in memory\n",
        "class make_dataset_multifield():\n",
        "\n",
        "    def __init__(self, mode, seed, fmaps, fparams, splits, fmaps_norm, verbose):\n",
        "\n",
        "        # get the total number of sims and maps\n",
        "        params_sims = np.loadtxt(fparams) #simulations parameters, NOT maps parameters\n",
        "        total_sims, total_maps, num_params = \\\n",
        "                params_sims.shape[0], params_sims.shape[0]*splits, params_sims.shape[1]\n",
        "        params_maps = np.zeros((total_maps, num_params), dtype=np.float32)\n",
        "        print(total_sims, splits, params_sims.shape[0])\n",
        "        for i in range(total_sims):\n",
        "            for j in range(splits):\n",
        "                params_maps[i*splits + j] = params_sims[i]\n",
        "\n",
        "        # normalize the value of the cosmological & astrophysical parameters\n",
        "        minimum     = np.array([0.1, 0.6, 0.25, 0.25, 0.5, 0.5])\n",
        "        maximum     = np.array([0.5, 1.0, 4.00, 4.00, 2.0, 2.0])\n",
        "        params_maps = (params_maps - minimum)/(maximum - minimum)\n",
        "\n",
        "        # get the size and offset depending on the type of dataset\n",
        "        if   mode=='train':  offset, size_sims = int(0.00*total_sims), int(0.90*total_sims)\n",
        "        elif mode=='valid':  offset, size_sims = int(0.90*total_sims), int(0.05*total_sims)\n",
        "        elif mode=='test':   offset, size_sims = int(0.95*total_sims), int(0.05*total_sims)\n",
        "        elif mode=='all':    offset, size_sims = int(0.00*total_sims), int(1.00*total_sims)\n",
        "        else:                raise Exception('Wrong name!')\n",
        "        size_maps = size_sims*splits\n",
        "\n",
        "        # randomly shuffle the simulations (not maps). Instead of 0 1 2 3...999 have a\n",
        "        # random permutation. E.g. 5 9 0 29...342\n",
        "        np.random.seed(seed)\n",
        "        sim_numbers = np.arange(total_sims) #shuffle sims not maps\n",
        "        np.random.shuffle(sim_numbers)\n",
        "        sim_numbers = sim_numbers[offset:offset+size_sims] #select indexes of mode\n",
        "\n",
        "        # get the corresponding indexes of the maps associated to the sims\n",
        "        indexes = np.zeros(size_maps, dtype=np.int32)\n",
        "        count = 0\n",
        "        for i in sim_numbers:\n",
        "            for j in range(splits):\n",
        "                indexes[count] = i*splits + j\n",
        "                count += 1\n",
        "\n",
        "        print(len(indexes))\n",
        "        print(indexes.shape)\n",
        "\n",
        "        # keep only the value of the parameters of the considered maps\n",
        "        params_maps = params_maps[indexes]\n",
        "\n",
        "        # define the matrix containing the maps with rotations and flipings\n",
        "        channels = len(fmaps)\n",
        "        dumb     = np.load(fmaps[0])    #[number of maps, height, width]\n",
        "        height, width = dumb.shape[1], dumb.shape[2];  del dumb\n",
        "        data     = np.zeros((size_maps*8, channels, height, width), dtype=np.float32)\n",
        "        params   = np.zeros((size_maps*8, num_params),              dtype=np.float32)\n",
        "\n",
        "        # read the data\n",
        "        print('Found %d channels\\nReading data...'%channels)\n",
        "        for channel, (fim, fnorm) in enumerate(zip(fmaps, fmaps_norm)):\n",
        "\n",
        "            # read maps in the considered channel\n",
        "            data_c = np.load(fim)\n",
        "            if data_c.shape[0]!=total_maps:  raise Exception('sizes do not match')\n",
        "            if verbose:  print('%.3e < F(all|orig) < %.3e'%(np.min(data_c), np.max(data_c)))\n",
        "\n",
        "            # rescale maps\n",
        "            if fim.find('Mstar')!=-1:  data_c = np.log10(data_c + 1.0)\n",
        "            else:                      data_c = np.log10(data_c)\n",
        "            if verbose:  print('%.3f < F(all|resc)  < %.3f'%(np.min(data_c), np.max(data_c)))\n",
        "\n",
        "            # normalize maps\n",
        "            if fnorm is None:\n",
        "                mean,    std     = np.mean(data_c), np.std(data_c)\n",
        "            else:\n",
        "                # read data\n",
        "                data_norm = np.load(fnorm)\n",
        "\n",
        "                # rescale\n",
        "                if fnorm.find('Mstar')!=-1:  data_norm = np.log10(data_norm + 1.0)\n",
        "                else:                        data_norm = np.log10(data_norm)\n",
        "\n",
        "                # compute mean and std\n",
        "                mean,    std     = np.mean(data_norm), np.std(data_norm)\n",
        "                del data_norm\n",
        "\n",
        "            data_c = (data_c - mean)/std\n",
        "            if verbose:  print('%.3f < F(all|norm) < %.3f'%(np.min(data_c), np.max(data_c)))\n",
        "\n",
        "            # keep only the data of the chosen set\n",
        "            data_c = data_c[indexes]\n",
        "\n",
        "            # do a loop over all rotations (each is 90 deg)\n",
        "            counted_maps = 0\n",
        "            for rot in [0,1,2,3]:\n",
        "                data_rot = np.rot90(data_c, k=rot, axes=(1,2))\n",
        "\n",
        "                data[counted_maps:counted_maps+size_maps,channel,:,:] = data_rot\n",
        "                params[counted_maps:counted_maps+size_maps]           = params_maps\n",
        "                counted_maps += size_maps\n",
        "\n",
        "                data[counted_maps:counted_maps+size_maps,channel,:,:] = \\\n",
        "                                                    np.flip(data_rot, axis=1)\n",
        "                params[counted_maps:counted_maps+size_maps]           = params_maps\n",
        "                counted_maps += size_maps\n",
        "\n",
        "            if verbose:\n",
        "                print('Channel %d contains %d maps'%(channel,counted_maps))\n",
        "                print('%.3f < F < %.3f\\n'%(np.min(data_c), np.max(data_c)))\n",
        "\n",
        "        self.size = data.shape[0]\n",
        "        self.x    = torch.tensor(data,   dtype=torch.float32)\n",
        "        self.y    = torch.tensor(params, dtype=torch.float32)\n",
        "        del data, data_c\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.size\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        return self.x[idx], self.y[idx]\n",
        "\n",
        "# This class creates the dataset. It will read the maps and store them in memory.\n",
        "# The rotations and flipings are done when calling the data\n",
        "class make_dataset_multifield2():\n",
        "\n",
        "    def __init__(self, mode, seed, fmaps, fparams, splits, fmaps_norm, verbose):\n",
        "\n",
        "        # get the total number of simulations and maps\n",
        "        params_sims = np.loadtxt(fparams) #simulations parameters, NOT maps parameters\n",
        "        total_sims, total_maps, num_params = \\\n",
        "                params_sims.shape[0], params_sims.shape[0]*splits, params_sims.shape[1]\n",
        "        params = np.zeros((total_maps, num_params), dtype=np.float32)\n",
        "        for i in range(total_sims):\n",
        "            for j in range(splits):\n",
        "                params[i*splits + j] = params_sims[i]\n",
        "\n",
        "        # normalize params\n",
        "        minimum = np.array([0.1, 0.6, 0.25, 0.25, 0.5, 0.5])\n",
        "        maximum = np.array([0.5, 1.0, 4.00, 4.00, 2.0, 2.0])\n",
        "        params  = (params - minimum)/(maximum - minimum)\n",
        "\n",
        "        # get the size and offset depending on the type of dataset\n",
        "        if   mode=='train':  offset, size_sims = int(0.00*total_sims), int(0.90*total_sims)\n",
        "        elif mode=='valid':  offset, size_sims = int(0.90*total_sims), int(0.05*total_sims)\n",
        "        elif mode=='test':   offset, size_sims = int(0.95*total_sims), int(0.05*total_sims)\n",
        "        elif mode=='all':    offset, size_sims = int(0.00*total_sims), int(1.00*total_sims)\n",
        "        else:                raise Exception('Wrong name!')\n",
        "        size_maps = size_sims*splits\n",
        "\n",
        "        # randomly shuffle the simulations (not maps). Instead of 0 1 2 3...999 have a\n",
        "        # random permutation. E.g. 5 9 0 29...342\n",
        "        np.random.seed(seed)\n",
        "        sim_numbers = np.arange(total_sims) #shuffle maps not rotations\n",
        "        np.random.shuffle(sim_numbers)\n",
        "        sim_numbers = sim_numbers[offset:offset+size_sims] #select indexes of mode\n",
        "\n",
        "        # get the corresponding indexes of the maps associated to the sims\n",
        "        indexes = np.zeros(size_maps, dtype=np.int32)\n",
        "        count = 0\n",
        "        for i in sim_numbers:\n",
        "            for j in range(splits):\n",
        "                indexes[count] = i*splits + j\n",
        "                count += 1\n",
        "\n",
        "        # keep only the value of the parameters of the considered maps\n",
        "        params = params[indexes]\n",
        "\n",
        "        # define the matrix containing the maps without rotations or flippings\n",
        "        channels = len(fmaps)\n",
        "        dumb     = np.load(fmaps[0])    #[number of maps, height, width]\n",
        "        height, width = dumb.shape[1], dumb.shape[2];  del dumb\n",
        "        data     = np.zeros((size_maps, channels, height, width), dtype=np.float32)\n",
        "\n",
        "        # read the data\n",
        "        print('Found %d channels\\nReading data...'%channels)\n",
        "        for channel, (fim, fnorm) in enumerate(zip(fmaps, fmaps_norm)):\n",
        "\n",
        "            # read maps in the considered channel\n",
        "            data_c = np.load(fim)\n",
        "            print(data_c.shape, total_maps)\n",
        "            if data_c.shape[0]!=total_maps:  raise Exception('sizes do not match')\n",
        "            if verbose:\n",
        "                print('%.3e < F(all|orig) < %.3e'%(np.min(data_c), np.max(data_c)))\n",
        "\n",
        "            # rescale maps\n",
        "            if fim.find('Mstar')!=-1:  data_c = np.log10(data_c + 1.0)\n",
        "            else:                      data_c = np.log10(data_c)\n",
        "            if verbose:\n",
        "                print('%.3f < F(all|resc)  < %.3f'%(np.min(data_c), np.max(data_c)))\n",
        "\n",
        "            # normalize maps\n",
        "            if fnorm is None:\n",
        "                mean,    std     = np.mean(data_c), np.std(data_c)\n",
        "                minimum, maximum = np.min(data_c),  np.max(data_c)\n",
        "            else:\n",
        "                # read data\n",
        "                data_norm     = np.load(fnorm)\n",
        "\n",
        "                # rescale data\n",
        "                if fnorm.find('Mstar')!=-1:  data_norm = np.log10(data_norm + 1.0)\n",
        "                else:                        data_norm = np.log10(data_norm)\n",
        "\n",
        "                # compute mean and std\n",
        "                mean,    std     = np.mean(data_norm), np.std(data_norm)\n",
        "                minimum, maximum = np.min(data_norm),  np.max(data_norm)\n",
        "                del data_norm\n",
        "\n",
        "            data_c = (data_c - mean)/std\n",
        "            if verbose:  print('%.3f < F(all|norm) < %.3f'%(np.min(data_c), np.max(data_c)))\n",
        "\n",
        "            # keep only the data of the chosen set\n",
        "            data[:,channel,:,:] = data_c[indexes]\n",
        "\n",
        "        self.size = data.shape[0]\n",
        "        self.x    = torch.tensor(data,   dtype=torch.float32)\n",
        "        self.y    = torch.tensor(params, dtype=torch.float32)\n",
        "        del data, data_c\n",
        "\n",
        "    def __len__(self):\n",
        "        return self.size\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "\n",
        "        # choose a rotation angle (0-0, 1-90, 2-180, 3-270)\n",
        "        # and whether do flipping or not\n",
        "        rot  = np.random.randint(0,4)\n",
        "        flip = np.random.randint(0,1)\n",
        "\n",
        "        # rotate and flip the maps\n",
        "        maps = torch.rot90(self.x[idx], k=rot, dims=[1,2])\n",
        "        if flip==1:  maps = torch.flip(maps, dims=[1])\n",
        "\n",
        "        return maps, self.y[idx]\n",
        "\n",
        "class model_o3_err(nn.Module):\n",
        "    def __init__(self, hidden, dr, channels):\n",
        "        super(model_o3_err, self).__init__()\n",
        "\n",
        "        # input: 1x256x256 ---------------> output: 2*hiddenx128x128\n",
        "        self.C01 = nn.Conv2d(channels,  2*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C02 = nn.Conv2d(2*hidden,  2*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C03 = nn.Conv2d(2*hidden,  2*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B01 = nn.BatchNorm2d(2*hidden)\n",
        "        self.B02 = nn.BatchNorm2d(2*hidden)\n",
        "        self.B03 = nn.BatchNorm2d(2*hidden)\n",
        "\n",
        "        # input: 2*hiddenx128x128 ----------> output: 4*hiddenx64x64\n",
        "        self.C11 = nn.Conv2d(2*hidden, 4*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C12 = nn.Conv2d(4*hidden, 4*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C13 = nn.Conv2d(4*hidden, 4*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B11 = nn.BatchNorm2d(4*hidden)\n",
        "        self.B12 = nn.BatchNorm2d(4*hidden)\n",
        "        self.B13 = nn.BatchNorm2d(4*hidden)\n",
        "\n",
        "        # input: 4*hiddenx64x64 --------> output: 8*hiddenx32x32\n",
        "        self.C21 = nn.Conv2d(4*hidden, 8*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C22 = nn.Conv2d(8*hidden, 8*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C23 = nn.Conv2d(8*hidden, 8*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B21 = nn.BatchNorm2d(8*hidden)\n",
        "        self.B22 = nn.BatchNorm2d(8*hidden)\n",
        "        self.B23 = nn.BatchNorm2d(8*hidden)\n",
        "\n",
        "        # input: 8*hiddenx32x32 ----------> output: 16*hiddenx16x16\n",
        "        self.C31 = nn.Conv2d(8*hidden,  16*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C32 = nn.Conv2d(16*hidden, 16*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C33 = nn.Conv2d(16*hidden, 16*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B31 = nn.BatchNorm2d(16*hidden)\n",
        "        self.B32 = nn.BatchNorm2d(16*hidden)\n",
        "        self.B33 = nn.BatchNorm2d(16*hidden)\n",
        "\n",
        "        # input: 16*hiddenx16x16 ----------> output: 32*hiddenx8x8\n",
        "        self.C41 = nn.Conv2d(16*hidden, 32*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C42 = nn.Conv2d(32*hidden, 32*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C43 = nn.Conv2d(32*hidden, 32*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B41 = nn.BatchNorm2d(32*hidden)\n",
        "        self.B42 = nn.BatchNorm2d(32*hidden)\n",
        "        self.B43 = nn.BatchNorm2d(32*hidden)\n",
        "\n",
        "        # input: 32*hiddenx8x8 ----------> output:64*hiddenx4x4\n",
        "        self.C51 = nn.Conv2d(32*hidden, 64*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C52 = nn.Conv2d(64*hidden, 64*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C53 = nn.Conv2d(64*hidden, 64*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B51 = nn.BatchNorm2d(64*hidden)\n",
        "        self.B52 = nn.BatchNorm2d(64*hidden)\n",
        "        self.B53 = nn.BatchNorm2d(64*hidden)\n",
        "\n",
        "        # input: 64*hiddenx4x4 ----------> output: 128*hiddenx1x1\n",
        "        self.C61 = nn.Conv2d(64*hidden, 128*hidden, kernel_size=4, stride=1, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B61 = nn.BatchNorm2d(128*hidden)\n",
        "\n",
        "        # self.P0  = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
        "\n",
        "        self.FC1  = nn.Linear(128*hidden, 64*hidden)\n",
        "        self.FC2  = nn.Linear(64*hidden,  12)\n",
        "\n",
        "        self.dropout   = nn.Dropout(p=dr)\n",
        "        # self.ReLU      = nn.ReLU()\n",
        "\n",
        "        self.LeakyReLUs = []\n",
        "\n",
        "        # We do this so that we can turn off some leaky relus according to our wish.\n",
        "        self.LeakyReLU1 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU2 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU3 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU4 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU5 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU6 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU7 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU8 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU9 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU10 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU11 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU12 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU13 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU14 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU15 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU16 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU17 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU18 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU19 = nn.LeakyReLU(0.2)\n",
        "        self.LeakyReLU20 = nn.LeakyReLU(0.2)\n",
        "\n",
        "        # self.tanh      = nn.Tanh()\n",
        "\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):\n",
        "                nn.init.constant_(m.weight, 1)\n",
        "                nn.init.constant_(m.bias, 0)\n",
        "            elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):\n",
        "                nn.init.kaiming_normal_(m.weight)\n",
        "\n",
        "    def forward(self, image):\n",
        "        x01 = self.C01(image)\n",
        "        x01_ = self.LeakyReLU1(x01)\n",
        "        x02 = self.C02(x01_)\n",
        "        x02_ = self.LeakyReLU2(self.B02(x02))\n",
        "        x03 = self.C03(x02_)\n",
        "        x03_ = self.LeakyReLU3(self.B03(x03))\n",
        "\n",
        "        x11 = self.C11(x03_)\n",
        "        x11_ = self.LeakyReLU4(self.B11(x11))\n",
        "        x12 = self.C12(x11_)\n",
        "        x12_ = self.LeakyReLU5(self.B12(x12))\n",
        "        x13 = self.C13(x12_)\n",
        "        x13_ = self.LeakyReLU6(self.B13(x13))\n",
        "\n",
        "        x21 = self.C21(x13_)\n",
        "        x21_ = self.LeakyReLU7(self.B21(x21))\n",
        "        x22 = self.C22(x21_)\n",
        "        x22_ = self.LeakyReLU8(self.B22(x22))\n",
        "        x23 = self.C23(x22_)\n",
        "        x23_ = self.LeakyReLU9(self.B23(x23))\n",
        "\n",
        "        x31 = self.C31(x23_)\n",
        "        x31_ = self.LeakyReLU10(self.B31(x31))\n",
        "        x32 = self.C32(x31_)\n",
        "        x32_ = self.LeakyReLU11(self.B32(x32))\n",
        "        x33 = self.C33(x32_)\n",
        "        x33_ = self.LeakyReLU12(self.B33(x33))\n",
        "\n",
        "        x41 = self.C41(x33_)\n",
        "        x41_ = self.LeakyReLU13(self.B41(x41))\n",
        "        x42 = self.C42(x41_)\n",
        "        x42_ = self.LeakyReLU14(self.B42(x42))\n",
        "        x43 = self.C43(x42_)\n",
        "        x43_ = self.LeakyReLU15(self.B43(x43))\n",
        "\n",
        "        x51 = self.C51(x43_)\n",
        "        x51_ = self.LeakyReLU16(self.B51(x51))\n",
        "        x52 = self.C52(x51_)\n",
        "        x52_ = self.LeakyReLU17(self.B52(x52))\n",
        "        x53 = self.C53(x52_)\n",
        "        x53_ = self.LeakyReLU18(self.B53(x53))\n",
        "\n",
        "        x61 = self.C61(x53_)\n",
        "        x61_ = self.LeakyReLU19(self.B61(x61))\n",
        "\n",
        "        x61__ = x61_.view(image.shape[0],-1)\n",
        "        x61__ = self.dropout(x61__)\n",
        "\n",
        "        x71 = self.FC1(x61__)\n",
        "        x71_ = self.LeakyReLU20(x71)\n",
        "        x71__ = self.dropout(x71_)\n",
        "        x81 = self.FC2(x71__)\n",
        "\n",
        "        # enforce the errors to be positive\n",
        "        y = torch.clone(x81)\n",
        "        y[:,6:12] = torch.square(x81[:,6:12])\n",
        "\n",
        "        return y, [x01_, x02_, x03_, x11_, x12_, x13_, x21_, x22_, x23_, x31_, x32_, x33_, x41_, x42_, x43_, x51_, x52_, x53_, x61_, x71_, x81]\n",
        "\n",
        "\n",
        "class model_o3_err_layer_output(nn.Module):\n",
        "    def __init__(self, hidden, dr, channels, layer_index=0):\n",
        "        super(model_o3_err_layer_output, self).__init__()\n",
        "\n",
        "        self.layer_index = layer_index\n",
        "\n",
        "        # input: 1x256x256 ---------------> output: 2*hiddenx128x128\n",
        "        self.C01 = nn.Conv2d(channels,  2*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C02 = nn.Conv2d(2*hidden,  2*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C03 = nn.Conv2d(2*hidden,  2*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B01 = nn.BatchNorm2d(2*hidden)\n",
        "        self.B02 = nn.BatchNorm2d(2*hidden)\n",
        "        self.B03 = nn.BatchNorm2d(2*hidden)\n",
        "\n",
        "        # input: 2*hiddenx128x128 ----------> output: 4*hiddenx64x64\n",
        "        self.C11 = nn.Conv2d(2*hidden, 4*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C12 = nn.Conv2d(4*hidden, 4*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C13 = nn.Conv2d(4*hidden, 4*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B11 = nn.BatchNorm2d(4*hidden)\n",
        "        self.B12 = nn.BatchNorm2d(4*hidden)\n",
        "        self.B13 = nn.BatchNorm2d(4*hidden)\n",
        "\n",
        "        # input: 4*hiddenx64x64 --------> output: 8*hiddenx32x32\n",
        "        self.C21 = nn.Conv2d(4*hidden, 8*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C22 = nn.Conv2d(8*hidden, 8*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C23 = nn.Conv2d(8*hidden, 8*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B21 = nn.BatchNorm2d(8*hidden)\n",
        "        self.B22 = nn.BatchNorm2d(8*hidden)\n",
        "        self.B23 = nn.BatchNorm2d(8*hidden)\n",
        "\n",
        "        # input: 8*hiddenx32x32 ----------> output: 16*hiddenx16x16\n",
        "        self.C31 = nn.Conv2d(8*hidden,  16*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C32 = nn.Conv2d(16*hidden, 16*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C33 = nn.Conv2d(16*hidden, 16*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B31 = nn.BatchNorm2d(16*hidden)\n",
        "        self.B32 = nn.BatchNorm2d(16*hidden)\n",
        "        self.B33 = nn.BatchNorm2d(16*hidden)\n",
        "\n",
        "        # input: 16*hiddenx16x16 ----------> output: 32*hiddenx8x8\n",
        "        self.C41 = nn.Conv2d(16*hidden, 32*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C42 = nn.Conv2d(32*hidden, 32*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C43 = nn.Conv2d(32*hidden, 32*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B41 = nn.BatchNorm2d(32*hidden)\n",
        "        self.B42 = nn.BatchNorm2d(32*hidden)\n",
        "        self.B43 = nn.BatchNorm2d(32*hidden)\n",
        "\n",
        "        # input: 32*hiddenx8x8 ----------> output:64*hiddenx4x4\n",
        "        self.C51 = nn.Conv2d(32*hidden, 64*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C52 = nn.Conv2d(64*hidden, 64*hidden, kernel_size=3, stride=1, padding=1,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.C53 = nn.Conv2d(64*hidden, 64*hidden, kernel_size=2, stride=2, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B51 = nn.BatchNorm2d(64*hidden)\n",
        "        self.B52 = nn.BatchNorm2d(64*hidden)\n",
        "        self.B53 = nn.BatchNorm2d(64*hidden)\n",
        "\n",
        "        # input: 64*hiddenx4x4 ----------> output: 128*hiddenx1x1\n",
        "        self.C61 = nn.Conv2d(64*hidden, 128*hidden, kernel_size=4, stride=1, padding=0,\n",
        "                            padding_mode='circular', bias=True)\n",
        "        self.B61 = nn.BatchNorm2d(128*hidden)\n",
        "\n",
        "        self.P0  = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)\n",
        "\n",
        "        self.FC1  = nn.Linear(128*hidden, 64*hidden)\n",
        "        self.FC2  = nn.Linear(64*hidden,  12)\n",
        "\n",
        "        self.dropout   = nn.Dropout(p=dr)\n",
        "        self.ReLU      = nn.ReLU()\n",
        "        self.LeakyReLU = nn.LeakyReLU(0.2)\n",
        "        self.tanh      = nn.Tanh()\n",
        "\n",
        "        for m in self.modules():\n",
        "            if isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):\n",
        "                nn.init.constant_(m.weight, 1)\n",
        "                nn.init.constant_(m.bias, 0)\n",
        "            elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Linear):\n",
        "                nn.init.kaiming_normal_(m.weight)\n",
        "\n",
        "\n",
        "    def forward(self, image):\n",
        "\n",
        "        x01 = self.C01(image)\n",
        "        x01_ = self.LeakyReLU(x01)\n",
        "        if self.layer_index == 0:\n",
        "          return x01_\n",
        "        x02 = self.C02(x01_)\n",
        "        x02_ = self.LeakyReLU(self.B02(x02))\n",
        "        if self.layer_index == 1:\n",
        "          return x02_\n",
        "        x03 = self.C03(x02_)\n",
        "        x03_ = self.LeakyReLU(self.B03(x03))\n",
        "        if self.layer_index == 2:\n",
        "          return x03_\n",
        "\n",
        "        x11 = self.C11(x03_)\n",
        "        x11_ = self.LeakyReLU(self.B11(x11))\n",
        "        if self.layer_index == 3:\n",
        "          return x11_\n",
        "        x12 = self.C12(x11_)\n",
        "        x12_ = self.LeakyReLU(self.B12(x12))\n",
        "        if self.layer_index == 4:\n",
        "          return x12_\n",
        "        x13 = self.C13(x12_)\n",
        "        x13_ = self.LeakyReLU(self.B13(x13))\n",
        "        if self.layer_index == 5:\n",
        "          return x13_\n",
        "\n",
        "        x21 = self.C21(x13_)\n",
        "        x21_ = self.LeakyReLU(self.B21(x21))\n",
        "        if self.layer_index == 6:\n",
        "          return x21_\n",
        "        x22 = self.C22(x21_)\n",
        "        x22_ = self.LeakyReLU(self.B22(x22))\n",
        "        if self.layer_index == 7:\n",
        "          return x22_\n",
        "        x23 = self.C23(x22_)\n",
        "        x23_ = self.LeakyReLU(self.B23(x23))\n",
        "        if self.layer_index == 8:\n",
        "          return x23_\n",
        "\n",
        "        x31 = self.C31(x23_)\n",
        "        x31_ = self.LeakyReLU(self.B31(x31))\n",
        "        if self.layer_index == 9:\n",
        "          return x31_\n",
        "        x32 = self.C32(x31_)\n",
        "        x32_ = self.LeakyReLU(self.B32(x32))\n",
        "        if self.layer_index == 10:\n",
        "          return x32_\n",
        "        x33 = self.C33(x32_)\n",
        "        x33_ = self.LeakyReLU(self.B33(x33))\n",
        "        if self.layer_index == 11:\n",
        "          return x33_\n",
        "\n",
        "        x41 = self.C41(x33_)\n",
        "        x41_ = self.LeakyReLU(self.B41(x41))\n",
        "        if self.layer_index == 12:\n",
        "          return x41_\n",
        "        x42 = self.C42(x41_)\n",
        "        x42_ = self.LeakyReLU(self.B42(x42))\n",
        "        if self.layer_index == 13:\n",
        "          return x42_\n",
        "        x43 = self.C43(x42_)\n",
        "        x43_ = self.LeakyReLU(self.B43(x43))\n",
        "        if self.layer_index == 14:\n",
        "          return x43_\n",
        "\n",
        "        x51 = self.C51(x43_)\n",
        "        x51_ = self.LeakyReLU(self.B51(x51))\n",
        "        if self.layer_index == 15:\n",
        "          return x51_\n",
        "        x52 = self.C52(x51_)\n",
        "        x52_ = self.LeakyReLU(self.B52(x52))\n",
        "        if self.layer_index == 16:\n",
        "          return x52_\n",
        "        x53 = self.C53(x52_)\n",
        "        x53_ = self.LeakyReLU(self.B53(x53))\n",
        "        if self.layer_index == 17:\n",
        "          return x53_\n",
        "\n",
        "        x61 = self.C61(x53_)\n",
        "        x61_ = self.LeakyReLU(self.B61(x61))\n",
        "        if self.layer_index == 18:\n",
        "          return x61_\n",
        "\n",
        "        x61__ = x61_.view(image.shape[0],-1)\n",
        "        x61__ = self.dropout(x61__)\n",
        "\n",
        "        x71 = self.FC1(x61__)\n",
        "        x71_ = self.LeakyReLU(x71)\n",
        "        if self.layer_index == 19:\n",
        "          return x71_\n",
        "        x71__ = self.dropout(x71_)\n",
        "        x81 = self.FC2(x71__)\n",
        "\n",
        "        # enforce the errors to be positive\n",
        "        y = torch.clone(x81)\n",
        "        y[:,6:12] = torch.square(x81[:,6:12])\n",
        "\n",
        "        if self.layer_index == 20:\n",
        "          return y\n",
        "\n",
        "\n",
        "def load_two_models_for_similarity_measure(fweights_1, fweights_2, fdatabase1, study_name, trial_number_1, trial_number_2, fdatabase2=None, layer_output=False, layer_index=0):\n",
        "  if fdatabase2 is None:\n",
        "    print('Assuming same database for the two models...')\n",
        "\n",
        "  if torch.cuda.is_available():\n",
        "    print(\"GPUs available\")\n",
        "    device = torch.device('cuda')\n",
        "  else:\n",
        "    print('GPUs not available')\n",
        "    device = torch.device('cpu')\n",
        "\n",
        "  fweights = [fweights_1, fweights_2]\n",
        "  fdatabases = [fdatabase1, fdatabase2]\n",
        "  trial_numbers = [trial_number_1, trial_number_2]\n",
        "  models = []\n",
        "\n",
        "  for i in range(2):\n",
        "    if fdatabase2 is None:\n",
        "      study = optuna.load_study(study_name=study_name, storage=fdatabase1)\n",
        "    else:\n",
        "      study = optuna.load_study(study_name=study_name, storage=fdatabases[i])\n",
        "\n",
        "    print(f\"\\nTrial number: {trial_numbers[i]}\")\n",
        "    trial_number = trial_numbers[i]\n",
        "    trial = study.trials[trial_number]\n",
        "    print(\"Trial number:  number {}\".format(trial.number))\n",
        "    print(\"Loss:          %.5e\"%trial.value)\n",
        "    print(\"Params: \")\n",
        "    for key, value in trial.params.items():\n",
        "      print(\"    {}: {}\".format(key, value))\n",
        "\n",
        "    if layer_output:\n",
        "      model = model_o3_err_layer_output(trial.params['hidden'], trial.params['dr'], 1, layer_index=layer_index)\n",
        "    else:\n",
        "      model = model_o3_err(trial.params['hidden'], trial.params['dr'], 1)\n",
        "    model = nn.DataParallel(model)\n",
        "    model.to(device=device)\n",
        "    network_total_params = sum(p.numel() for p in model.parameters())\n",
        "    print('total number of parameters in the model = %d'%network_total_params)\n",
        "\n",
        "    if os.path.exists(fweights[i]):\n",
        "      model.load_state_dict(torch.load(fweights[i], map_location=torch.device(device)))\n",
        "      print('Weights loaded')\n",
        "      models.append(model)\n",
        "    else:\n",
        "      raise Exception('file with weights not found!!!')\n",
        "\n",
        "  if len(models) != 2:\n",
        "    print(\"Warning! Two models were not loaded...\")\n",
        "  return models\n",
        "\n",
        "# def load_two_models_for_similarity_measure_modified(fweights_1, fweights_2, fdatabase1, study_name, trial_number_1, trial_number_2, fdatabase2=None):  # modified to also return linear layer.\n",
        "#   if fdatabase2 is None:\n",
        "#     print('Assuming same database for the two models...')\n",
        "\n",
        "#   if torch.cuda.is_available():\n",
        "#     print(\"GPUs available\")\n",
        "#     device = torch.device('cuda')\n",
        "#   else:\n",
        "#     print('GPUs not available')\n",
        "#     device = torch.device('cpu')\n",
        "\n",
        "#   fweights = [fweights_1, fweights_2]\n",
        "#   fdatabases = [fdatabase1, fdatabase2]\n",
        "#   trial_numbers = [trial_number_1, trial_number_2]\n",
        "#   models = []\n",
        "#   lr_models = []\n",
        "\n",
        "#   for i in range(2):\n",
        "#     if fdatabase2 is None:\n",
        "#       study = optuna.load_study(study_name=study_name, storage=fdatabase1)\n",
        "#     else:\n",
        "#       study = optuna.load_study(study_name=study_name, storage=fdatabases[i])\n",
        "\n",
        "#     print(f\"\\nTrial number: {trial_numbers[i]}\")\n",
        "#     trial_number = trial_numbers[i]\n",
        "#     trial = study.trials[trial_number]\n",
        "#     print(\"Trial number:  number {}\".format(trial.number))\n",
        "#     print(\"Loss:          %.5e\"%trial.value)\n",
        "#     print(\"Params: \")\n",
        "#     for key, value in trial.params.items():\n",
        "#       print(\"    {}: {}\".format(key, value))\n",
        "\n",
        "#     model = model_o3_err(trial.params['hidden'], trial.params['dr'], 1)\n",
        "#     model = nn.DataParallel(model)\n",
        "#     model.to(device=device)\n",
        "#     network_total_params = sum(p.numel() for p in model.parameters())\n",
        "#     print('total number of parameters in the model = %d'%network_total_params)\n",
        "\n",
        "#     if os.path.exists(fweights[i]):\n",
        "#       model.load_state_dict(torch.load(fweights[i], map_location=torch.device(device)))\n",
        "#       print('Weights loaded')\n",
        "#       models.append(model)\n",
        "#     else:\n",
        "#       raise Exception('file with weights not found!!!')\n",
        "\n",
        "#     lr_models.append(\n",
        "#         linear_reg()\n",
        "#     )\n",
        "\n",
        "#   if len(models) != 2:\n",
        "#     print(\"Warning! Two models were not loaded...\")\n",
        "#   if len(lr_models) != 2:\n",
        "#     print(\"Warning! Two linear layer models were not loaded...\")\n",
        "#   return models, lr_models\n",
        "\n",
        "import seaborn as sns\n",
        "sns.set_context(\"paper\", font_scale = 2)\n",
        "sns.set_style('whitegrid')\n",
        "sns.set(style='ticks')\n",
        "def plot_cka(sim):\n",
        "  fig, ax = plt.subplots(1, 1, figsize=(5, 5))\n",
        "\n",
        "  divider = make_axes_locatable(ax)\n",
        "  cax = divider.append_axes('right', size='5%', pad=0.05)\n",
        "\n",
        "  im = ax.imshow(sim, cmap='magma', vmin=0.0,vmax=1.0)\n",
        "  layers = np.arange(1, 21, 2)\n",
        "  ax.set_xticks(layers); ax.set_yticks(layers)\n",
        "  ax.set_xticklabels(layers); ax.set_yticklabels(layers)\n",
        "  ax.axes.invert_yaxis()\n",
        "  fig.colorbar(im, cax=cax, orientation='vertical')\n",
        "  return ax\n",
        "\n",
        "\n",
        "def get_test_acc(model, test_loader):\n",
        "  # get the number of maps in the test set\n",
        "  num_maps = 0\n",
        "  for x,y in test_loader:\n",
        "        num_maps += x.shape[0]\n",
        "  print('\\nNumber of maps in the test set: %d'%num_maps)\n",
        "\n",
        "  # define the arrays containing the value of the parameters\n",
        "  params_true = np.zeros((num_maps,6), dtype=np.float32)\n",
        "  params_NN   = np.zeros((num_maps,6), dtype=np.float32)\n",
        "  errors_NN   = np.zeros((num_maps,6), dtype=np.float32)\n",
        "\n",
        "  # get test loss\n",
        "  test_loss1, test_loss2 = torch.zeros(len(g)).to(device), torch.zeros(len(g)).to(device)\n",
        "  test_loss, points = 0.0, 0\n",
        "  model.eval()\n",
        "  for x, y in test_loader:\n",
        "      with torch.no_grad():\n",
        "          bs    = x.shape[0]    #batch size\n",
        "          x     = x.to(device)  #send data to device\n",
        "          y     = y.to(device)  #send data to device\n",
        "          p     = model(x)[0]      #prediction for mean and variance  # EDIT: Take 0th element because we modify the architecture to return the features as well.\n",
        "          y_NN  = p[:,:6]       #prediction for mean\n",
        "          e_NN  = p[:,6:]       #prediction for error\n",
        "          loss1 = torch.mean((y_NN[:,g] - y[:,g])**2,                     axis=0)\n",
        "          loss2 = torch.mean(((y_NN[:,g] - y[:,g])**2 - e_NN[:,g]**2)**2, axis=0)\n",
        "          test_loss1 += loss1*bs\n",
        "          test_loss2 += loss2*bs\n",
        "\n",
        "          # save results to their corresponding arrays\n",
        "          params_true[points:points+x.shape[0]] = y.cpu().numpy()\n",
        "          params_NN[points:points+x.shape[0]]   = y_NN.cpu().numpy()\n",
        "          errors_NN[points:points+x.shape[0]]   = e_NN.cpu().numpy()\n",
        "          points    += x.shape[0]\n",
        "  test_loss = torch.log(test_loss1/points) + torch.log(test_loss2/points)\n",
        "  test_loss = torch.mean(test_loss).item()\n",
        "  print('Test loss = %.3e\\n'%test_loss)\n",
        "  return params_true, params_NN, errors_NN\n",
        "\n",
        "def get_test_acc_modified(model_list, test_loader):  # for internal layer accuracy experiment\n",
        "  lr_model = model_list[1]\n",
        "  model = model_list[0]\n",
        "  # get the number of maps in the test set\n",
        "  num_maps = 0\n",
        "  for x,y in test_loader:\n",
        "        num_maps += x.shape[0]\n",
        "  print('\\nNumber of maps in the test set: %d'%num_maps)\n",
        "\n",
        "  # define the arrays containing the value of the parameters\n",
        "  params_true = np.zeros((num_maps,6), dtype=np.float32)\n",
        "  params_NN   = np.zeros((num_maps,6), dtype=np.float32)\n",
        "  errors_NN   = np.zeros((num_maps,6), dtype=np.float32)\n",
        "\n",
        "  # get test loss\n",
        "  test_loss1, test_loss2 = torch.zeros(len(g)).to(device), torch.zeros(len(g)).to(device)\n",
        "  test_loss, points = 0.0, 0\n",
        "  model.eval(); lr_model.eval()\n",
        "  for x, y in test_loader:\n",
        "      with torch.no_grad():\n",
        "          bs    = x.shape[0]    #batch size\n",
        "          x     = x.to(device)  #send data to device\n",
        "          y     = y.to(device)  #send data to device\n",
        "          pp   = model(x)[0]           #NN output\n",
        "          p    = lr_model(pp)  # Logistic regression.\n",
        "          y_NN  = p[:,:6]       #prediction for mean\n",
        "          e_NN  = p[:,6:]       #prediction for error\n",
        "          loss1 = torch.mean((y_NN[:,g] - y[:,g])**2,                     axis=0)\n",
        "          loss2 = torch.mean(((y_NN[:,g] - y[:,g])**2 - e_NN[:,g]**2)**2, axis=0)\n",
        "          test_loss1 += loss1*bs\n",
        "          test_loss2 += loss2*bs\n",
        "\n",
        "          # save results to their corresponding arrays\n",
        "          params_true[points:points+x.shape[0]] = y.cpu().numpy()\n",
        "          params_NN[points:points+x.shape[0]]   = y_NN.cpu().numpy()\n",
        "          errors_NN[points:points+x.shape[0]]   = e_NN.cpu().numpy()\n",
        "          points    += x.shape[0]\n",
        "  test_loss = torch.log(test_loss1/points) + torch.log(test_loss2/points)\n",
        "  test_loss = torch.mean(test_loss).item()\n",
        "  print('Test loss = %.3e\\n'%test_loss)\n",
        "  return params_true, params_NN, errors_NN\n",
        "\n",
        "from sklearn.metrics import r2_score, mean_squared_error\n",
        "\n",
        "def get_r2_score(params_true, params_NN):\n",
        "  r2_scores_params = []\n",
        "  for i in range(params_true.shape[1]):\n",
        "    r2_scores_params.append(\n",
        "      r2_score(params_true[:, i], params_NN[:, i])\n",
        "    )\n",
        "  return r2_scores_params\n",
        "\n",
        "def get_rmse(params_true, params_NN):\n",
        "  rmse_scores_params = []\n",
        "  for i in range(params_true.shape[1]):\n",
        "    rmse_scores_params.append(\n",
        "      mean_squared_error(params_true[:, i], params_NN[:, i])\n",
        "    )\n",
        "  return rmse_scores_params\n",
        "\n",
        "class linear_reg(nn.Module):\n",
        "  def __init__(self):\n",
        "    super(linear_reg, self).__init__()\n",
        "    self.linear_reg = nn.LazyLinear(12)\n",
        "\n",
        "    # for m in self.modules():\n",
        "    #   if isinstance(m, nn.LazyLinear):\n",
        "    #     nn.init.kaiming_normal_(m.weight)\n",
        "\n",
        "  def forward(self, x):\n",
        "      x = self.linear_reg(x)\n",
        "      # enforce the errors to be positive\n",
        "      y = torch.clone(x)\n",
        "      y[:,6:12] = torch.square(x[:,6:12])\n",
        "      return y\n",
        "\n",
        "# class model_o3_err_C01(nn.Module):\n",
        "#   def __init__(self, hidden, dr, channels, state_dict):\n",
        "#     super(model_o3_err_C01, self).__init__()\n",
        "#       # input: 1x256x256 ---------------> output: 2*hiddenx128x128\n",
        "#     self.C01 = nn.Conv2d(channels, 2*hidden, kernel_size=3, stride=1, padding=1,\n",
        "#                             padding_mode='circular', bias=True)\n",
        "\n",
        "#     with torch.no_grad():\n",
        "#       self.C01.weight.copy_(state_dict['module.C01.weight'])\n",
        "#       self.C01.bias.copy_(state_dict['module.C01.bias'])\n",
        "\n",
        "#     self.C01.weight.requires_grad = False\n",
        "#     self.C01.bias.requires_grad = False\n",
        "\n",
        "#     self.LeakyReLU = nn.LeakyReLU(0.2)\n",
        "\n",
        "#   def forward(self, image):\n",
        "#     x01 = self.C01(image)\n",
        "#     x01_ = self.LeakyReLU(x01)\n",
        "#     return x01_\n",
        "\n",
        "def perform_training(model_list, train_loader, valid_loader, epochs):\n",
        "  model, lr_model = model_list\n",
        "  optimizer = torch.optim.AdamW(lr_model.parameters(), lr=lr, weight_decay=wd, betas=(beta1, beta2))\n",
        "  scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.3, patience=10)\n",
        "\n",
        "  print('Computing initial validation loss')\n",
        "  lr_model.eval(); model.eval()\n",
        "  valid_loss1, valid_loss2 = torch.zeros(len(g)).to(device), torch.zeros(len(g)).to(device)\n",
        "  min_valid_loss, points = 0.0, 0\n",
        "  for x, y in valid_loader:\n",
        "        with torch.no_grad():\n",
        "            bs   = x.shape[0]                #batch size\n",
        "            x    = x.to(device=device)       #maps\n",
        "            y    = y.to(device=device)[:,g]  #parameters\n",
        "\n",
        "            #with torch.no_grad():  # No weights must be updated. Only the linear regressor in the next line is trained.\n",
        "            pp   = model(x)[0]           #NN output\n",
        "            p    = lr_model(pp)  # Logistic regression.\n",
        "\n",
        "            y_NN = p[:,g]                    #posterior mean\n",
        "            e_NN = p[:,h]                    #posterior std\n",
        "            loss1 = torch.mean((y_NN - y)**2,                axis=0)\n",
        "            loss2 = torch.mean(((y_NN - y)**2 - e_NN**2)**2, axis=0)\n",
        "            loss  = torch.mean(torch.log(loss1) + torch.log(loss2))\n",
        "            valid_loss1 += loss1*bs\n",
        "            valid_loss2 += loss2*bs\n",
        "            points += bs\n",
        "  min_valid_loss = torch.log(valid_loss1/points) + torch.log(valid_loss2/points)\n",
        "  min_valid_loss = torch.mean(min_valid_loss).item()\n",
        "  print('Initial valid loss = %.3e'%min_valid_loss)\n",
        "\n",
        "  lr_model.train(); model.eval()\n",
        "  # do a loop over all epochs\n",
        "  start = time.time()\n",
        "  for epoch in range(epochs):\n",
        "      # do training\n",
        "      train_loss1, train_loss2 = torch.zeros(len(g)).to(device), torch.zeros(len(g)).to(device)\n",
        "      train_loss, points = 0.0, 0\n",
        "      lr_model.train(); model.eval()\n",
        "      for x, y in train_loader:\n",
        "          bs   = x.shape[0]         #batch size\n",
        "          x    = x.to(device)       #maps\n",
        "          y    = y.to(device)[:,g]  #parameters\n",
        "\n",
        "          #with torch.no_grad():  # No weights must be updated. Only the logistic regressor in the next line is trained.\n",
        "          pp   = model(x)[0]           #NN output\n",
        "          p    = lr_model(pp)  # Logistic regression.\n",
        "\n",
        "          y_NN = p[:,g]             #posterior mean\n",
        "          e_NN = p[:,h]             #posterior std\n",
        "          loss1 = torch.mean((y_NN - y)**2,                axis=0)\n",
        "          loss2 = torch.mean(((y_NN - y)**2 - e_NN**2)**2, axis=0)\n",
        "          loss  = torch.mean(torch.log(loss1) + torch.log(loss2))\n",
        "          train_loss1 += loss1*bs\n",
        "          train_loss2 += loss2*bs\n",
        "          points      += bs\n",
        "          optimizer.zero_grad()\n",
        "          loss.backward()\n",
        "          optimizer.step()\n",
        "\n",
        "          #if points>18000:  break\n",
        "      train_loss = torch.log(train_loss1/points) + torch.log(train_loss2/points)\n",
        "      train_loss = torch.mean(train_loss).item()\n",
        "\n",
        "      # do validation: cosmo alone & all params\n",
        "      valid_loss1, valid_loss2 = torch.zeros(len(g)).to(device), torch.zeros(len(g)).to(device)\n",
        "      valid_loss, points = 0.0, 0\n",
        "      lr_model.eval(); model.eval()\n",
        "      for x, y in valid_loader:\n",
        "          with torch.no_grad():\n",
        "              bs    = x.shape[0]         #batch size\n",
        "              x     = x.to(device)       #maps\n",
        "              y     = y.to(device)[:,g]  #parameters\n",
        "\n",
        "              pp    = model(x)[0]           #NN output\n",
        "              p     = lr_model(pp)  # Logistic regression.\n",
        "\n",
        "              y_NN  = p[:,g]             #posterior mean\n",
        "              e_NN  = p[:,h]             #posterior std\n",
        "              loss1 = torch.mean((y_NN - y)**2,                axis=0)\n",
        "              loss2 = torch.mean(((y_NN - y)**2 - e_NN**2)**2, axis=0)\n",
        "              loss  = torch.mean(torch.log(loss1) + torch.log(loss2))\n",
        "              valid_loss1 += loss1*bs\n",
        "              valid_loss2 += loss2*bs\n",
        "              points     += bs\n",
        "      valid_loss = torch.log(valid_loss1/points) + torch.log(valid_loss2/points)\n",
        "      valid_loss = torch.mean(valid_loss).item()\n",
        "\n",
        "      scheduler.step(valid_loss)\n",
        "\n",
        "      # verbose\n",
        "      print('%03d %.3e %.3e '%(epoch, train_loss, valid_loss), end='')\n",
        "\n",
        "      # save model if it is better\n",
        "      if valid_loss<min_valid_loss:\n",
        "          torch.save(lr_model.state_dict(), fmodel)\n",
        "          min_valid_loss = valid_loss\n",
        "          print('(C) ', end='')\n",
        "      print('')\n",
        "\n",
        "      # save losses to file\n",
        "      f = open(floss, 'a')\n",
        "      f.write('%d %.5e %.5e\\n'%(epoch, train_loss, valid_loss))\n",
        "      f.close()\n",
        "\n",
        "  stop = time.time()\n",
        "  print('Time take (h):', \"{:.4f}\".format((stop-start)/3600.0))\n",
        "\n",
        "  return lr_model"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fdX-SDYSMDOg",
        "outputId": "789ac4dd-77a0-47f2-85a8-2c3bb65ea8e8"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting optuna==2.4.0\n",
            "  Downloading optuna-2.4.0-py3-none-any.whl (282 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m282.7/282.7 kB\u001b[0m \u001b[31m3.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting alembic (from optuna==2.4.0)\n",
            "  Downloading alembic-1.11.2-py3-none-any.whl (225 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m225.3/225.3 kB\u001b[0m \u001b[31m6.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting cliff (from optuna==2.4.0)\n",
            "  Downloading cliff-4.3.0-py3-none-any.whl (80 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m80.6/80.6 kB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hCollecting cmaes>=0.6.0 (from optuna==2.4.0)\n",
            "  Downloading cmaes-0.10.0-py3-none-any.whl (29 kB)\n",
            "Collecting colorlog (from optuna==2.4.0)\n",
            "  Downloading colorlog-6.7.0-py2.py3-none-any.whl (11 kB)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.10/dist-packages (from optuna==2.4.0) (1.3.1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.10/dist-packages (from optuna==2.4.0) (1.22.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from optuna==2.4.0) (23.1)\n",
            "Requirement already satisfied: scipy!=1.4.0 in /usr/local/lib/python3.10/dist-packages (from optuna==2.4.0) (1.10.1)\n",
            "Requirement already satisfied: sqlalchemy>=1.1.0 in /usr/local/lib/python3.10/dist-packages (from optuna==2.4.0) (2.0.19)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from optuna==2.4.0) (4.65.0)\n",
            "Requirement already satisfied: typing-extensions>=4.2.0 in /usr/local/lib/python3.10/dist-packages (from sqlalchemy>=1.1.0->optuna==2.4.0) (4.7.1)\n",
            "Requirement already satisfied: greenlet!=0.4.17 in /usr/local/lib/python3.10/dist-packages (from sqlalchemy>=1.1.0->optuna==2.4.0) (2.0.2)\n",
            "Collecting Mako (from alembic->optuna==2.4.0)\n",
            "  Downloading Mako-1.2.4-py3-none-any.whl (78 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m78.7/78.7 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: PrettyTable>=0.7.2 in /usr/local/lib/python3.10/dist-packages (from cliff->optuna==2.4.0) (0.7.2)\n",
            "Requirement already satisfied: PyYAML>=3.12 in /usr/local/lib/python3.10/dist-packages (from cliff->optuna==2.4.0) (6.0.1)\n",
            "Collecting autopage>=0.4.0 (from cliff->optuna==2.4.0)\n",
            "  Downloading autopage-0.5.1-py3-none-any.whl (29 kB)\n",
            "Collecting cmd2>=1.0.0 (from cliff->optuna==2.4.0)\n",
            "  Downloading cmd2-2.4.3-py3-none-any.whl (147 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m147.2/147.2 kB\u001b[0m \u001b[31m7.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: importlib-metadata>=4.4 in /usr/lib/python3/dist-packages (from cliff->optuna==2.4.0) (4.6.4)\n",
            "Collecting stevedore>=2.0.1 (from cliff->optuna==2.4.0)\n",
            "  Downloading stevedore-5.1.0-py3-none-any.whl (49 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.6/49.6 kB\u001b[0m \u001b[31m5.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: attrs>=16.3.0 in /usr/local/lib/python3.10/dist-packages (from cmd2>=1.0.0->cliff->optuna==2.4.0) (23.1.0)\n",
            "Collecting pyperclip>=1.6 (from cmd2>=1.0.0->cliff->optuna==2.4.0)\n",
            "  Downloading pyperclip-1.8.2.tar.gz (20 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Requirement already satisfied: wcwidth>=0.1.7 in /usr/local/lib/python3.10/dist-packages (from cmd2>=1.0.0->cliff->optuna==2.4.0) (0.2.6)\n",
            "Collecting pbr!=2.1.0,>=2.0.0 (from stevedore>=2.0.1->cliff->optuna==2.4.0)\n",
            "  Downloading pbr-5.11.1-py2.py3-none-any.whl (112 kB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m112.7/112.7 kB\u001b[0m \u001b[31m9.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: MarkupSafe>=0.9.2 in /usr/local/lib/python3.10/dist-packages (from Mako->alembic->optuna==2.4.0) (2.1.3)\n",
            "Building wheels for collected packages: pyperclip\n",
            "  Building wheel for pyperclip (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for pyperclip: filename=pyperclip-1.8.2-py3-none-any.whl size=11123 sha256=ed7339b41413c5e946ab96d37ca8da3a8bc82377c4f64312603bc2844a3c75cd\n",
            "  Stored in directory: /root/.cache/pip/wheels/04/24/fe/140a94a7f1036003ede94579e6b4227fe96c840c6f4dcbe307\n",
            "Successfully built pyperclip\n",
            "Installing collected packages: pyperclip, pbr, Mako, colorlog, cmd2, cmaes, autopage, stevedore, alembic, cliff, optuna\n",
            "Successfully installed Mako-1.2.4 alembic-1.11.2 autopage-0.5.1 cliff-4.3.0 cmaes-0.10.0 cmd2-2.4.3 colorlog-6.7.0 optuna-2.4.0 pbr-5.11.1 pyperclip-1.8.2 stevedore-5.1.0\n",
            "CUDA Not Available\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2. Analysis start"
      ],
      "metadata": {
        "id": "dn7Jeya5MHC_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "base_dir = '/content/drive/MyDrive/CAMELS/After_ICML_results'\n",
        "field = 'T'\n",
        "data_sim = 'IllustrisTNG'\n",
        "model_sim = 'IllustrisTNG'\n",
        "study_name = 'wd_dr_hidden_lr_o3'\n",
        "\n",
        "# data parameters\n",
        "fmaps      = ['/content/maps_T.npy'] #tuple containing the maps with the different fields to consider\n",
        "fmaps_norm = [None if data_sim == model_sim else os.path.join(base_dir, f'Maps_{field}_Nbody_{model_sim}_LH_z=0.00.npy' if field == 'Mtot' else f'Maps_{field}_{model_sim}_LH_z=0.00.npy')] #if you want to normalize the maps according to the properties of some data set, put that data set here (This is mostly used when training on IllustrisTNG and testing on SIMBA, or vicerversa)\n",
        "fparams    = os.path.join(base_dir, f'params_LH_Nbody_{data_sim}.txt' if field == 'Mtot' else f'params_LH_{data_sim}.txt')  # Note: f'params_LH_Nbody_{data_sim}.txt' ONLY for Nbody, else f'params_LH_{data_sim}.txt'\n",
        "seed       = 1   #random seed to split maps among training, validation and testing\n",
        "splits     = 1   #number of maps per simulation\n",
        "\n",
        "channels        = 1                #we only consider here 1 field\n",
        "params          = [0,1,2,3,4,5]    #0(Omega_m) 1(sigma_8) 2(A_SN1) 3 (A_AGN1) 4(A_SN2) 5(A_AGN2). The code will be trained to predict all these parameters.\n",
        "g               = params           #g will contain the mean of the posterior\n",
        "h               = [6+i for i in g] #h will contain the variance of the posterior\n",
        "rot_flip_in_mem = False            #whether rotations and flipings are kept in memory. True will make the code faster but consumes more RAM memory.\n",
        "\n",
        "fmaps2 = os.path.join(base_dir, f'Maps_{field}_Nbody_{data_sim}_LH_z=0.00.npy' if field == 'Mtot' else f'Maps_{field}_{data_sim}_LH_z=0.00.npy')\n",
        "maps  = np.load(fmaps2)\n",
        "print('Shape of the maps:',maps.shape)\n",
        "\n",
        "# define the array that will contain the indexes of the maps\n",
        "indexes = np.zeros(1000*splits, dtype=np.int32)\n",
        "\n",
        "# do a loop over all maps and choose the ones we want\n",
        "count = 0\n",
        "for i in range(15000):\n",
        "    if i%15 in np.arange(splits):\n",
        "      indexes[count] = i\n",
        "      count += 1\n",
        "print('Selected %d maps out of 15000'%count)\n",
        "\n",
        "# save these maps to a new file\n",
        "maps = maps[indexes]\n",
        "np.save('maps_T.npy', maps)\n",
        "print(f'Shape of selected maps array: {maps.shape}')\n",
        "del maps\n",
        "\n",
        "# IMPORTANT: shuffle must be False to ensure iteration over train loader is stable.\n",
        "# We need the iterations to be stable since after saving, the outputs will be used for training.\n",
        "batch_size = 1\n",
        "train_loader = create_dataset_multifield('train', seed, fmaps, fparams, batch_size, splits, fmaps_norm,\n",
        "                                         rot_flip_in_mem=rot_flip_in_mem, verbose=True, shuffle=False)\n",
        "test_loader  = create_dataset_multifield('test', seed, fmaps, fparams, batch_size, splits, fmaps_norm,\n",
        "                                        rot_flip_in_mem=False, verbose=True, shuffle=False)\n",
        "\n",
        "true_params = []\n",
        "for i, (x, y) in enumerate(train_loader):\n",
        "    with torch.no_grad():\n",
        "      bs   = x.shape[0]         #batch size\n",
        "      x    = x.to(device)       #maps\n",
        "      y    = y.to(device)[:, g]  #parameters\n",
        "      true_params.append(y)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gct_ep22Lw5Z",
        "outputId": "6785f417-633f-40a5-b794-71e5002d370f"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Shape of the maps: (15000, 256, 256)\n",
            "Selected 1000 maps out of 15000\n",
            "Shape of selected maps array: (1000, 256, 256)\n",
            "Found 1 channels\n",
            "Reading data...\n",
            "(1000, 256, 256) 1000\n",
            "1.547e+03 < F(all|orig) < 8.150e+07\n",
            "3.189 < F(all|resc)  < 7.911\n",
            "-1.265 < F(all|norm) < 4.516\n",
            "Found 1 channels\n",
            "Reading data...\n",
            "(1000, 256, 256) 1000\n",
            "1.547e+03 < F(all|orig) < 8.150e+07\n",
            "3.189 < F(all|resc)  < 7.911\n",
            "-1.265 < F(all|norm) < 4.516\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "todo: we need to use chunking fit only for layer_index=0,1,2....so modify code.\n",
        "todo: test R2 on test set not train set."
      ],
      "metadata": {
        "id": "ptFPik_BZSGh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def chunker(seq, size):\n",
        "  return (seq[pos:pos + size] for pos in range(0, len(seq), size))\n",
        "\n",
        "def std_dev(array, mean, n):  # Use this instead of np.std since we cannot load the entire 900 features from earlier layers into memory at once.\n",
        "  return np.sqrt((1/n) * np.sum((array-mean)**2))\n",
        "\n",
        "param_index = 0  # First parameter only\n",
        "layer_index = 18\n",
        "BATCH_SIZE = 10  # For training regressor below\n",
        "trial_num = 10  # or 28 for ICML quick experiment.\n",
        "\n",
        "import glob\n",
        "if layer_index <= 2:  # memory issues if entire array is loaded in memory. So read in 900 files, one for each example in the training set.\n",
        "  files = f'/content/drive/MyDrive/CAMELS/After_ICML_results/layer_outputs_for_trials_used_for_ICML/{field}_{model_sim}_{data_sim}_trial{trial_num}_all900outputs_layerIndex{layer_index}batchnum*.pt'  # note: this path should ideally also include trial number info...\n",
        "  assert len(files) == 900\n",
        "  sorted_files = sorted(glob.glob(files), key = lambda x: int(x.split('/')[8].split('_')[6][8:].split('.')[0]))\n",
        "  def yield_data_splits():\n",
        "    for f, par in zip(\n",
        "        chunker(sorted_files, BATCH_SIZE),\n",
        "        chunker(true_params, BATCH_SIZE)\n",
        "    ):\n",
        "      t = torch.vstack([torch.load(ff) for ff in f])\n",
        "      params = [p[0][param_index] for p in par]\n",
        "      yield t, params, t.mean()\n",
        "elif layer_index >= 3:\n",
        "  files = f'/content/drive/MyDrive/CAMELS/After_ICML_results/layer_outputs_for_trials_used_for_ICML/{field}_{model_sim}_{data_sim}_trial{trial_num}_all900outputs_layerIndex{layer_index}.pt'  # note: this path should ideally also include trial number info...\n",
        "  assert len(glob.glob(files)) == 1\n",
        "  outputs = torch.load(files)\n",
        "  assert outputs.shape[0] == 900\n",
        "  def yield_data_splits():\n",
        "    for t, par in zip(\n",
        "        chunker(outputs, BATCH_SIZE),\n",
        "        chunker(true_params, BATCH_SIZE)\n",
        "    ):\n",
        "      params = [p[0][param_index] for p in par]\n",
        "      yield t, params, t.mean()\n",
        "\n",
        "if layer_index <= 2:\n",
        "  _t = next(yield_data_splits())\n",
        "  t_shape = torch.numel(_t[0]) * (len(sorted_files) / BATCH_SIZE)\n",
        "  mean_all = np.sum([ds[2]*torch.numel(ds[0]) for ds in yield_data_splits()]) / t_shape\n",
        "elif layer_index >= 3:\n",
        "  mean_all = outputs.mean()\n",
        "\n",
        "if layer_index <= 2:\n",
        "  def yield_meand_sums():\n",
        "    for f, par in zip(\n",
        "        chunker(sorted_files, BATCH_SIZE),\n",
        "        chunker(true_params, BATCH_SIZE)\n",
        "    ):\n",
        "      t = torch.vstack([torch.load(ff) for ff in f])\n",
        "      params = [p[0][param_index] for p in par]\n",
        "\n",
        "      yield torch.sum((t - mean_all)**2)\n",
        "\n",
        "  len_meand_sums = sum(1 for _ in yield_meand_sums())\n",
        "  meand_sums = np.fromiter((ms for ms in yield_meand_sums()), dtype=float)\n",
        "\n",
        "  assert len_meand_sums == (len(sorted_files) / BATCH_SIZE)\n",
        "elif layer_index >= 3:\n",
        "  def yield_meand_sums():\n",
        "    for t, par in zip(\n",
        "        chunker(outputs, BATCH_SIZE),\n",
        "        chunker(true_params, BATCH_SIZE)\n",
        "    ):\n",
        "      params = [p[0][param_index] for p in par]\n",
        "\n",
        "      yield torch.sum((t - mean_all)**2)\n",
        "\n",
        "if layer_index <= 2:\n",
        "  stddev_all = np.sqrt((1/t_shape) * np.sum(meand_sums))  # For formula, see https://en.wikipedia.org/wiki/Standard_deviation#Relationship_between_standard_deviation_and_mean\n",
        "else:\n",
        "  stddev_all = outputs.std()\n",
        "\n",
        "if layer_index >= 3:\n",
        "  _t = next(yield_data_splits())\n",
        "  t_shape = torch.numel(_t[0]) * (900 / BATCH_SIZE)\n",
        "  mean_all_check = np.sum([ds[2]*torch.numel(ds[0]) for ds in yield_data_splits()]) / t_shape\n",
        "  assert np.allclose(\n",
        "      mean_all_check,\n",
        "      outputs.mean()\n",
        "  )\n",
        "  len_meand_sums = sum(1 for _ in yield_meand_sums())\n",
        "  meand_sums = np.fromiter((ms for ms in yield_meand_sums()), dtype=float)\n",
        "  std_all_check = np.sqrt((1/t_shape) * np.sum(meand_sums))\n",
        "  assert np.allclose(\n",
        "      std_all_check,\n",
        "      outputs.std()\n",
        "  )\n",
        "print(outputs.min(), outputs.max(), mean_all, stddev_all)"
      ],
      "metadata": {
        "id": "Th9A_xvcVY7L",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a37205f7-4d24-4bb3-adc5-aa724d7e1f7a"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "tensor(-0.0859) tensor(0.2103) tensor(0.0024) tensor(0.0265)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "todo: check this code works for small test cases..."
      ],
      "metadata": {
        "id": "FDhmUSFRDGZ7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "stddev_all, mean_all"
      ],
      "metadata": {
        "id": "uMr5qQjT8Pqw",
        "outputId": "4db35f04-3c88-4998-9d8b-46a6f32fbb9d",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(tensor(0.0265), tensor(0.0024))"
            ]
          },
          "metadata": {},
          "execution_count": 4
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.linear_model import Ridge, LinearRegression, SGDRegressor\n",
        "from sklearn.preprocessing import StandardScaler\n",
        "\n",
        "reg = SGDRegressor()\n",
        "for _ in range(10): # 10 passes through the data\n",
        "  for X, y, _ in yield_data_splits():\n",
        "    X = (X - mean_all) / stddev_all\n",
        "    X = X.reshape(X.shape[0], -1)\n",
        "    reg.partial_fit(X, y)"
      ],
      "metadata": {
        "id": "-iT5fGPfVQK6"
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "y_preds = []\n",
        "y_true = []\n",
        "for X, y, _ in yield_data_splits():\n",
        "  X = (X - mean_all) / stddev_all\n",
        "  X = X.reshape(X.shape[0], -1)\n",
        "  y_pred = reg.predict(X)\n",
        "  y_preds.append(y_pred)\n",
        "  y_true.append(y)\n",
        "\n",
        "r2s = r2_score(np.hstack(y_true), np.hstack(y_preds))\n",
        "print(f'R2 score with layer_index {layer_index}, trial {trial_num} = {r2s}')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "7MqAB-nH32Ih",
        "outputId": "4d584036-2b38-4848-8d03-5feedde11e0a"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "R2 score with layer_index 18, trial 10 = -5.3750752205003535e+20\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import gc\n",
        "def get_outputs(model, train_loader, prefix=None, layer_index=0):  # index specifies output of which layer we want to extract.\n",
        "  model.eval()\n",
        "  for i, (x, y) in enumerate(train_loader):\n",
        "    with torch.no_grad():\n",
        "      bs   = x.shape[0]         #batch size\n",
        "      x    = x.to(device)       #maps\n",
        "      y    = y.to(device)[:, g]  #parameters\n",
        "\n",
        "      out  = model(x)           #NN output\n",
        "      if layer_index == 0 or layer_index == 1 or layer_index == 2:  # memory issues, so we save each output.\n",
        "        # torch.save(out, f'{field}_{model_sim}_{data_sim}_trial{trial_number_1}_batchsize{batch_size}_layerIndex{layer_index}_batchnum{i}.pt')\n",
        "        del out\n",
        "        gc.collect()\n",
        "      else:\n",
        "        yield out\n",
        "\n",
        "def yield_data_splits_for_test():\n",
        "  for t, par in zip(\n",
        "      chunker(outputs, BATCH_SIZE),\n",
        "      chunker(true_params, BATCH_SIZE)\n",
        "  ):\n",
        "    params = [p[0][param_index] for p in par]\n",
        "    yield t, params\n",
        "\n",
        "true_params = []\n",
        "for i, (_, y) in enumerate(test_loader):\n",
        "    with torch.no_grad():\n",
        "      bs   = x.shape[0]         #batch size\n",
        "      # x    = x.to(device)       #maps\n",
        "      y    = y.to(device)[:,g]  #parameters\n",
        "      true_params.append(y)\n",
        "\n",
        "# trials_1 = np.load(os.path.join(base_dir, 'trials_1.npy'))\n",
        "# trials_2 = np.load(os.path.join(base_dir, 'trials_2.npy'))\n",
        "\n",
        "# assert len(np.intersect1d(trials_1, trials_2)) == 0\n",
        "\n",
        "fweights_1 = os.path.join(base_dir, f'weights_{model_sim}_{field}_{trial_num}_all_steps_500_500_o3.pt')\n",
        "fdatabase1 = 'sqlite:////' + os.path.join(base_dir, f'{model_sim}_o3_{field}_all_steps_500_500_o3.db')\n",
        "\n",
        "study = optuna.load_study(study_name=study_name, storage=fdatabase1)\n",
        "\n",
        "trial_number = int(fweights_1.split('/')[6].split('_')[3])\n",
        "\n",
        "model = load_model_for_similarity_measure(fweights_1, fdatabase1, study_name, trial_number, layer_output=True, layer_index=layer_index)\n",
        "\n",
        "outputs = get_outputs(model, test_loader, layer_index=layer_index)\n",
        "outputs = torch.vstack([arr for arr in outputs])\n",
        "\n",
        "y_preds = []\n",
        "y_true = []\n",
        "for X, y in yield_data_splits_for_test():\n",
        "  X = X.reshape(X.shape[0], -1)\n",
        "  X = (X - mean_all) / stddev_all\n",
        "  y_pred = reg.predict(X)\n",
        "  y_preds.append(y_pred)\n",
        "  y_true.append(y)\n",
        "\n",
        "r2s = r2_score(np.hstack(y_true), np.hstack(y_preds))\n",
        "print(f'R2 score with layer_index {layer_index}, trial {trial_num} = {r2s}')"
      ],
      "metadata": {
        "id": "IIjQFiX2Udvm",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "7b8cabab-fd54-496e-9150-5d5fd5d67cea"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "GPUs not available\n",
            "\n",
            "Trial number: 10\n",
            "Trial number:  number 10\n",
            "Loss:          -1.22274e+01\n",
            "Params: \n",
            "    dr: 0.6695702848487869\n",
            "    hidden: 6\n",
            "    lr: 0.0011543514200712233\n",
            "    wd: 3.859932400788152e-07\n",
            "total number of parameters in the model = 8466876\n",
            "Weights loaded\n",
            "R2 score with layer_index 10, trial 10 = -1.5864497833485992e+28\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "todo: training happens for one trial but I am testing on two trials -- fix."
      ],
      "metadata": {
        "id": "XP8Q59PZNi40"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "r2_score(np.hstack(y_true), np.hstack(y_preds))  # # with layer_index=0"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "RN_fLNlnbfjR",
        "outputId": "82ecba79-b34d-407a-f71b-b3ad799b2400"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.21904098479058076"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "r2_score(np.hstack(y_true), np.hstack(y_preds))  # # with layer_index=20"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fZhJlNq_Q6gm",
        "outputId": "49b5ebf4-c0df-44aa-dbcf-0d2fae3c63cd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.8809224141966819"
            ]
          },
          "metadata": {},
          "execution_count": 28
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "r2_score(np.hstack(y_true), np.hstack(y_preds))  # # with layer_index=10"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kPvzUKAMSzve",
        "outputId": "fb1dea6f-6058-4ce0-e72b-aa1bb0598a11"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.38652268214278784"
            ]
          },
          "metadata": {},
          "execution_count": 23
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "r2_score(np.hstack(y_true), np.hstack(y_preds))  # # with layer_index=16"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OEIeQE0fa5-W",
        "outputId": "ce30f701-fb9c-4ad0-925f-727596d86a45"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.9613126182628591"
            ]
          },
          "metadata": {},
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "r2_score(np.hstack(y_true), np.hstack(y_preds))  # # with layer_index=18"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "52SluLfneosc",
        "outputId": "304f13b3-f58d-4c02-9bd6-eb1d8e1e4bc6"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "0.9753024655199403"
            ]
          },
          "metadata": {},
          "execution_count": 13
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "y_preds  # with layer_index=0"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3BWywnkpVhVD",
        "outputId": "3b23bb6a-0e5e-49b0-8ab6-106506193dde"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[array([-1.74921162e+14, -1.22941702e+15,  1.15113647e+15, -3.54342462e+14,\n",
              "         1.26445390e+14,  3.02976815e+14, -3.28154966e+14,  8.32266892e+14,\n",
              "         1.36467119e+15,  6.87493922e+14]),\n",
              " array([ 4.52157570e+14, -9.15411201e+14, -1.55280093e+14, -7.36774970e+14,\n",
              "        -1.87113722e+14, -2.56441961e+14,  1.26406706e+14, -2.12743179e+14,\n",
              "         1.13585875e+15,  1.30814196e+15]),\n",
              " array([-1.08765192e+15,  6.28530446e+14, -4.79874470e+14, -1.24491919e+15,\n",
              "         1.82060987e+14, -3.51950213e+14,  3.46235231e+14,  2.63217652e+14,\n",
              "        -2.65794876e+14,  9.40113987e+14]),\n",
              " array([-1.18786307e+15, -4.65894348e+14, -1.68375060e+14,  1.01303510e+15,\n",
              "         4.28621519e+13,  3.89659114e+14,  6.75878222e+14,  2.92869915e+14,\n",
              "         4.52482279e+14, -4.71503418e+14]),\n",
              " array([-8.09659679e+14, -2.73468372e+14,  5.92967614e+14, -5.51126195e+14,\n",
              "         3.96729200e+14,  1.20832523e+15,  5.85339633e+14, -1.09643273e+15,\n",
              "         2.54221537e+14,  8.68517490e+14]),\n",
              " array([ 3.45783870e+14, -1.70795288e+14, -1.41353425e+14, -7.72154610e+14,\n",
              "         1.26829064e+15,  1.67355894e+15, -8.27563555e+14,  1.13467520e+14,\n",
              "         1.11844195e+15, -3.95629478e+13]),\n",
              " array([-7.44652107e+14,  8.69826924e+14, -8.82656405e+14,  3.92161403e+14,\n",
              "        -2.06001145e+14,  4.97023290e+14,  2.00604587e+14, -3.00668837e+14,\n",
              "        -2.36919083e+14, -8.18832560e+13]),\n",
              " array([ 6.09942873e+14,  6.47103074e+14,  9.17624322e+14,  1.74242895e+14,\n",
              "         3.57545112e+14,  1.50032335e+15, -9.31834983e+14,  5.38644486e+14,\n",
              "         8.18520373e+14,  1.23164290e+15]),\n",
              " array([ 4.22952117e+14,  5.46738675e+12,  2.96687019e+14, -1.46195971e+15,\n",
              "         9.95870971e+14,  2.49787570e+14,  9.24764352e+14,  1.75574208e+14,\n",
              "        -3.06704881e+14, -5.67035928e+14]),\n",
              " array([ 4.75212675e+14,  6.24611511e+14,  3.81010477e+14,  4.16781127e+14,\n",
              "         5.91555632e+14,  1.11634539e+15, -7.47110278e+13, -1.15046149e+15,\n",
              "        -6.98625669e+13, -4.26582631e+13]),\n",
              " array([ 4.86542134e+14,  4.24592703e+12,  1.33791751e+14,  5.20151530e+14,\n",
              "         2.18168379e+14, -1.44146826e+15,  1.40051481e+13, -7.95412174e+14,\n",
              "        -1.67612658e+15,  7.85473980e+14]),\n",
              " array([-1.80390670e+15, -3.21382388e+14,  4.14255679e+14, -1.31468549e+15,\n",
              "         2.56711996e+14, -9.26010647e+14, -9.67665841e+14,  2.76374816e+14,\n",
              "         1.51420985e+15,  1.09617224e+15]),\n",
              " array([-9.65441596e+13,  2.83463247e+14,  1.08686403e+15, -1.31251646e+15,\n",
              "         6.62452220e+14,  3.94173894e+14,  7.08079909e+14, -4.42469570e+14,\n",
              "        -7.43112360e+13, -7.14154479e+14]),\n",
              " array([ 6.28100790e+14,  7.77283900e+14, -9.16927401e+14,  1.24501775e+15,\n",
              "        -7.70456948e+13,  1.11645191e+15,  5.24057678e+14,  1.96213604e+14,\n",
              "        -1.88121581e+15, -1.54317019e+14]),\n",
              " array([ 1.80176212e+15,  6.98892141e+14,  1.28640081e+15,  1.14006283e+14,\n",
              "         7.98281790e+14,  8.27946712e+14,  2.90541425e+14,  9.06039753e+14,\n",
              "        -8.06941366e+14, -1.27754799e+15]),\n",
              " array([-1.49184312e+15, -1.22429763e+13, -2.42299601e+14,  9.00448939e+14,\n",
              "        -3.88182740e+14,  7.62600741e+14,  2.37976641e+14, -4.37633190e+14,\n",
              "        -5.39389713e+14,  8.69860647e+13]),\n",
              " array([-3.33410759e+14,  7.38650696e+14, -8.66024289e+14,  2.59970969e+14,\n",
              "         2.47590096e+14, -7.71071973e+14, -2.41131012e+14, -3.50557540e+14,\n",
              "         1.11829367e+15,  5.60053571e+14]),\n",
              " array([-1.85855388e+14, -5.50308396e+14, -4.03111575e+14,  2.46211046e+13,\n",
              "        -7.40512719e+13,  1.07944154e+15, -9.32014425e+13, -4.57885328e+14,\n",
              "         1.52632822e+14, -1.01730471e+15]),\n",
              " array([ 7.08326846e+14,  3.16196105e+14, -2.08614973e+14,  1.00874365e+15,\n",
              "         6.04239504e+14, -2.32243510e+13,  7.47798338e+14,  4.93380132e+14,\n",
              "         2.66176967e+14, -7.66849124e+14]),\n",
              " array([-2.09953256e+15,  2.43416484e+14, -4.65889951e+14,  4.28961010e+14,\n",
              "         8.47338331e+14, -1.94288095e+14,  8.33817626e+12,  6.81655000e+14,\n",
              "         9.78482634e+14,  9.10092224e+14]),\n",
              " array([-6.49625437e+14, -5.28127257e+14,  2.93538742e+13, -5.65054224e+12,\n",
              "        -8.83057129e+14, -1.84207017e+14, -9.43693355e+13, -2.98651599e+14,\n",
              "        -1.06607943e+15,  1.13378105e+15]),\n",
              " array([ 7.01905233e+14,  8.65865179e+14,  5.29589922e+14,  1.81809851e+13,\n",
              "         5.54399657e+14, -6.06906272e+14, -9.74613348e+14,  9.14264126e+14,\n",
              "         3.30134161e+14, -2.17177220e+15]),\n",
              " array([-4.73401647e+14, -7.51344967e+14,  5.87358591e+14, -6.34975483e+13,\n",
              "         1.94985667e+15,  6.82734365e+14, -5.02751556e+14,  2.53187748e+14,\n",
              "         1.71065276e+15,  1.93006333e+14]),\n",
              " array([ 3.27588315e+14, -9.47750631e+14, -6.22339594e+14, -2.55754730e+14,\n",
              "        -1.45379761e+14,  2.99916493e+14,  3.93527346e+14, -1.18482146e+15,\n",
              "         2.27062965e+14,  2.05426391e+14]),\n",
              " array([ 1.22206289e+14, -5.30720255e+14, -2.82720083e+14, -3.32916283e+14,\n",
              "        -8.48418614e+14,  9.19802395e+13,  1.36760054e+14,  6.80661204e+14,\n",
              "        -5.05221326e+14, -5.33581600e+14]),\n",
              " array([ 1.21047341e+15, -1.52709776e+15, -1.55671761e+14, -9.49979187e+14,\n",
              "        -5.86897327e+13,  2.30188295e+14,  1.43688835e+15,  2.19246741e+13,\n",
              "        -2.56761393e+14,  7.72418978e+14]),\n",
              " array([-6.20752965e+14, -2.37785600e+13,  9.80257898e+14,  9.47808017e+13,\n",
              "        -3.95969701e+14, -1.13977017e+14,  3.34357226e+14, -1.69911554e+14,\n",
              "         1.01779054e+15, -3.23809180e+13]),\n",
              " array([-2.71473026e+14, -1.81129063e+13, -2.26917259e+14,  5.61694877e+14,\n",
              "        -1.01755206e+15,  4.69077159e+14,  5.29182190e+14, -3.88556878e+14,\n",
              "        -9.87429268e+14, -5.70252226e+14]),\n",
              " array([ 6.37569246e+14, -7.07125309e+13,  1.18569096e+15,  9.96609532e+14,\n",
              "        -2.33527948e+14,  2.92134266e+14, -1.31313721e+15,  9.12126642e+14,\n",
              "         1.65319759e+15, -3.13710417e+14]),\n",
              " array([-2.38507757e+14, -7.68517606e+14,  2.16708839e+14,  1.37433216e+15,\n",
              "         1.60699874e+15,  1.07772076e+15,  7.89070247e+13, -6.32324973e+14,\n",
              "        -3.31498066e+14,  4.53639797e+14]),\n",
              " array([-1.43627283e+15,  4.01264217e+14,  4.67406780e+14, -4.18157709e+14,\n",
              "         3.94202568e+14, -4.87828774e+14, -1.64690899e+15,  1.03865963e+15,\n",
              "         3.19638089e+14,  1.35617127e+15]),\n",
              " array([-3.53357050e+14,  9.33152881e+13,  9.53312607e+14,  4.90823918e+14,\n",
              "        -7.60329755e+14,  8.76093524e+14, -1.61956425e+15,  4.58100328e+13,\n",
              "         1.59792519e+14,  8.48153119e+14]),\n",
              " array([ 4.56662713e+13, -1.38165429e+14,  2.89994958e+14, -2.45815908e+14,\n",
              "        -7.75079306e+14,  6.73151293e+14, -8.30723084e+13, -1.63500389e+15,\n",
              "        -3.59704781e+14,  8.29473609e+14]),\n",
              " array([ 1.01621411e+15, -2.62009789e+13,  3.38683917e+13, -7.89450607e+14,\n",
              "        -2.21393252e+14,  3.42420782e+14, -1.23608719e+12,  6.04838373e+13,\n",
              "        -6.08691639e+14,  6.27180370e+13]),\n",
              " array([-7.44657738e+14,  1.50915643e+15,  1.16991311e+15,  1.28634258e+15,\n",
              "         8.62828439e+14,  9.15688006e+14,  2.10964378e+14,  1.70752814e+14,\n",
              "         3.97870863e+14,  1.64794599e+14]),\n",
              " array([-1.39515515e+15,  6.85937919e+14, -5.24966688e+14,  1.00073152e+15,\n",
              "        -2.99994867e+14, -4.52136215e+14,  1.90376819e+15, -1.63154282e+14,\n",
              "        -2.95692450e+14, -2.91721807e+14]),\n",
              " array([ 9.85766906e+14,  8.57656483e+14,  6.69730153e+13, -5.51004433e+14,\n",
              "        -3.33357529e+14,  1.24230245e+15,  2.88773866e+14,  6.77903521e+14,\n",
              "        -9.56608247e+14, -3.89617848e+12]),\n",
              " array([ 9.02190050e+14,  1.01355291e+15, -1.32894972e+14,  4.95769835e+14,\n",
              "         6.27443622e+14, -1.13911803e+15, -5.65103866e+14, -7.70277806e+14,\n",
              "         9.16027970e+14,  6.39942177e+13]),\n",
              " array([1.46827005e+15, 4.68562732e+14, 2.67960751e+14, 9.72724595e+14,\n",
              "        2.42058575e+14, 6.91674634e+13, 1.03348707e+15, 4.87664974e+14,\n",
              "        3.83637368e+14, 1.28345793e+15]),\n",
              " array([-1.13629107e+14, -8.01263883e+14,  7.18382267e+14, -7.34470504e+14,\n",
              "        -7.97097475e+13,  7.87022092e+14,  4.05113242e+12,  1.38101016e+15,\n",
              "         6.98849787e+14, -2.71137759e+14]),\n",
              " array([ 1.17660976e+15,  6.83736804e+14, -2.96965273e+14, -2.09730869e+14,\n",
              "         5.59088391e+14,  4.79769941e+14,  3.10977837e+14,  5.86759380e+14,\n",
              "         6.70415884e+14,  8.92761712e+14]),\n",
              " array([-3.91978025e+14, -4.84313202e+14,  7.11183686e+14,  3.65304039e+14,\n",
              "         2.33206731e+14, -6.16343715e+13,  1.82648286e+14,  7.43391400e+14,\n",
              "         5.84466526e+14,  6.96619542e+14]),\n",
              " array([ 3.62736140e+14, -3.05925988e+14,  1.57449040e+14, -3.44921193e+14,\n",
              "         1.45090844e+14,  2.91862730e+14, -1.37598096e+15,  1.35004699e+15,\n",
              "         9.91438203e+14, -8.00823221e+14]),\n",
              " array([-2.60974660e+14, -1.07222409e+15,  2.95524642e+14, -7.84914199e+13,\n",
              "         5.12280634e+14, -4.97381719e+14,  1.90023003e+15,  3.50981300e+14,\n",
              "         3.79460431e+14,  9.50531905e+14]),\n",
              " array([-1.46356436e+14, -4.43007256e+13, -1.45665022e+14, -1.44984704e+15,\n",
              "         7.62571097e+14,  1.02885562e+15,  6.73553313e+14,  1.03318111e+14,\n",
              "        -4.80400333e+14,  6.09238508e+14]),\n",
              " array([-2.77963964e+14,  8.80273757e+14,  3.43082520e+14,  1.37022278e+15,\n",
              "         4.99273397e+14,  2.17035096e+14,  4.74120722e+14,  1.08152572e+15,\n",
              "         1.93910183e+14, -9.08105939e+14]),\n",
              " array([ 6.39602427e+14, -8.59458090e+14, -5.78870292e+14, -2.56405926e+13,\n",
              "         1.03594659e+14,  4.33084997e+14,  3.44737624e+14,  1.48208962e+14,\n",
              "         1.77228493e+15,  7.65429018e+14]),\n",
              " array([-1.29552284e+15, -1.15624455e+15,  2.04728399e+14,  4.16563261e+14,\n",
              "         3.57396541e+14,  7.61741484e+14,  8.57947859e+14, -2.79368228e+14,\n",
              "        -5.50983840e+13,  3.10043335e+14]),\n",
              " array([-9.25092951e+14, -4.45076835e+14,  1.22376542e+15,  2.80981504e+14,\n",
              "         1.59980473e+14,  2.51972712e+13,  3.83772678e+14,  5.41666350e+14,\n",
              "         1.18418009e+15,  6.63388808e+14]),\n",
              " array([ 1.55041619e+14,  4.08036299e+14,  5.67060340e+14, -4.34482737e+13,\n",
              "         5.01244850e+14,  1.49928518e+15, -1.39664294e+14,  1.63408543e+14,\n",
              "         3.47084638e+14,  1.38859745e+15]),\n",
              " array([ 4.29475042e+14, -9.84154475e+14,  1.78751775e+14, -2.87184615e+14,\n",
              "         5.53691977e+14,  4.03406476e+14,  1.12432249e+15, -1.10709633e+15,\n",
              "        -2.52630642e+14, -3.00109606e+14]),\n",
              " array([ 1.85653793e+15, -3.89550799e+14,  3.13282458e+14,  5.89899395e+14,\n",
              "         3.65263970e+14, -1.07307178e+15,  6.38169504e+14, -3.21739625e+13,\n",
              "         5.89466798e+14,  1.68623793e+15]),\n",
              " array([ 9.38962111e+14, -4.50832087e+14, -8.45361621e+14, -4.94119046e+14,\n",
              "         3.96547059e+14, -5.34288015e+14,  1.13035966e+14,  1.15015043e+15,\n",
              "        -1.41522360e+14,  9.47493251e+14]),\n",
              " array([-1.34702816e+15,  4.03566916e+14,  1.51961700e+15, -9.18629850e+14,\n",
              "         3.87524946e+14, -8.25186151e+14, -8.72933242e+14,  5.81369937e+14,\n",
              "        -1.33171081e+15,  1.00285276e+14]),\n",
              " array([ 9.83012550e+14, -3.14095042e+14,  9.33680391e+14, -3.33194197e+14,\n",
              "        -7.44716298e+13, -1.17652876e+15, -1.08987932e+15,  1.10236905e+14,\n",
              "        -7.13494093e+13,  8.68394514e+14]),\n",
              " array([ 5.32439911e+14,  8.37040825e+14,  2.65758511e+14,  5.69104061e+14,\n",
              "         3.22301654e+14,  1.50751810e+15, -1.25579063e+15,  3.23258553e+14,\n",
              "        -1.11508096e+15,  1.08944401e+15]),\n",
              " array([ 2.08510921e+14, -6.44126376e+14,  4.16936535e+14, -6.63385515e+14,\n",
              "        -5.27571717e+14,  4.03840272e+14,  7.02858121e+14,  5.50746639e+14,\n",
              "         7.91397297e+14, -3.27309743e+14]),\n",
              " array([ 1.22367225e+14,  5.76514987e+14, -1.01961965e+15,  5.50080011e+14,\n",
              "        -3.52501975e+14,  4.09970358e+14,  1.21422930e+14,  9.05279550e+14,\n",
              "         1.03872506e+15,  7.34222953e+14]),\n",
              " array([-3.89748374e+14,  4.49813452e+14,  6.05804411e+14,  1.64194912e+15,\n",
              "        -1.35576233e+14, -1.76923822e+14, -7.83823656e+14, -6.94900451e+14,\n",
              "         1.05617010e+13, -1.01964059e+14]),\n",
              " array([-7.14465787e+14, -6.36237253e+14,  4.36490191e+14, -1.14382725e+15,\n",
              "         5.61739945e+14,  6.80101314e+14,  1.27166859e+15, -1.18314076e+15,\n",
              "         1.16012283e+15,  1.01638454e+15]),\n",
              " array([-1.20259325e+14, -1.03821720e+14, -1.34098076e+14,  7.75628017e+14,\n",
              "        -1.65397045e+15, -5.30734148e+14,  9.96469623e+14, -5.95700925e+14,\n",
              "         8.24728028e+13,  2.89295745e+14]),\n",
              " array([-1.82203241e+15, -9.49884420e+14,  3.72125322e+14, -1.90665018e+14,\n",
              "        -2.53047022e+14, -3.26689663e+13,  9.96483610e+13,  3.78701922e+14,\n",
              "         4.23939846e+14,  9.92653259e+14]),\n",
              " array([ 7.27772444e+14, -7.53262706e+14,  5.69439532e+14, -1.11559859e+13,\n",
              "         1.33647741e+15, -4.97090459e+14, -5.72014580e+14,  2.17048952e+14,\n",
              "         3.35391454e+14,  6.38295906e+14]),\n",
              " array([-1.46136904e+15,  6.38385195e+14,  2.47627415e+14,  2.24923922e+14,\n",
              "        -2.89025936e+14, -1.38541024e+15, -3.66624489e+14,  6.35571069e+14,\n",
              "        -1.93454313e+14,  8.17753656e+14]),\n",
              " array([-3.29860882e+14, -1.99751891e+14, -1.39441351e+15,  6.41183083e+14,\n",
              "         9.09594518e+14,  8.69290648e+14,  3.61296247e+14,  2.20888858e+14,\n",
              "         1.27620978e+13, -1.33526380e+15]),\n",
              " array([-1.95695638e+14,  5.32835515e+14, -2.96850841e+15, -9.40611784e+14,\n",
              "         1.95003818e+15, -1.33267045e+15,  1.54825224e+14,  5.42186315e+14,\n",
              "         3.72971885e+14,  8.78084634e+14]),\n",
              " array([-2.31202160e+14, -1.65467999e+15, -5.12809768e+13,  1.19910502e+15,\n",
              "         2.35606981e+14, -6.79724244e+14, -7.11496331e+14,  6.40532787e+14,\n",
              "         2.04367005e+14,  2.58035667e+14]),\n",
              " array([ 2.59002992e+14, -3.66427279e+14,  1.98710428e+14, -7.68623607e+14,\n",
              "         9.22739398e+13,  2.60137388e+14, -8.27173186e+14,  6.92610182e+14,\n",
              "         7.08802724e+14, -1.67408505e+15]),\n",
              " array([ 4.70212147e+13, -5.45020608e+14,  2.50000223e+14, -1.14626177e+15,\n",
              "         6.94779880e+14,  3.91290937e+14,  4.67354698e+14,  9.49345858e+14,\n",
              "        -1.67831512e+14,  2.80078917e+14]),\n",
              " array([ 5.45620043e+14,  7.28048454e+14,  3.97103821e+14,  5.04141391e+14,\n",
              "        -2.34486778e+14,  8.81475637e+14,  3.46572152e+14,  1.22571533e+15,\n",
              "        -2.75409152e+14,  3.23400877e+14]),\n",
              " array([-6.94933084e+14, -2.12342954e+15,  1.02665399e+15, -4.78229767e+14,\n",
              "         3.50643864e+13,  6.24874146e+14,  5.98194995e+14, -1.00667766e+15,\n",
              "         1.54886442e+14, -1.26755990e+14]),\n",
              " array([-8.71383128e+14,  4.66516611e+14,  5.01757152e+12,  5.11654219e+14,\n",
              "        -1.90500912e+14,  4.35198236e+14, -6.85585162e+13, -3.62134771e+14,\n",
              "         9.18106307e+14, -8.29047166e+14]),\n",
              " array([-1.27944720e+15, -1.41366427e+15,  5.16197833e+13,  9.60415237e+14,\n",
              "        -6.96066379e+14,  2.41940946e+14,  1.07638132e+14, -1.30211508e+15,\n",
              "         1.86567207e+15, -1.52454024e+15]),\n",
              " array([ 3.38036186e+14,  1.96917123e+14, -5.43449043e+13, -1.32716825e+15,\n",
              "        -1.50847590e+15,  1.78110410e+15, -1.70907212e+14,  1.77977767e+15,\n",
              "         5.85270952e+14, -1.15452955e+15]),\n",
              " array([ 7.52887065e+13,  3.59153857e+13,  6.22554875e+14, -1.27038049e+15,\n",
              "        -2.81702367e+14,  5.83283352e+14,  8.74221092e+14,  5.98908995e+14,\n",
              "        -1.07534680e+15, -8.12484678e+13]),\n",
              " array([-1.98351400e+15, -1.03942311e+14, -1.47975423e+15,  5.22272687e+14,\n",
              "        -3.29787099e+13,  3.11342624e+14,  8.38412922e+14, -3.96579203e+14,\n",
              "        -2.82046124e+15, -9.06558453e+14]),\n",
              " array([ 9.10655109e+14, -4.86485831e+14, -1.33921819e+15, -6.35050694e+14,\n",
              "         5.41791866e+13, -1.82301308e+15,  6.96175446e+14,  2.34036439e+14,\n",
              "        -8.57897807e+14, -4.55816617e+14]),\n",
              " array([ 1.07808303e+15, -8.74288370e+13,  6.30464355e+14,  2.34946402e+14,\n",
              "         2.94159200e+14, -7.93444858e+14, -1.46979645e+15,  9.38314991e+14,\n",
              "        -5.29656012e+14, -7.61005648e+14]),\n",
              " array([ 2.38798738e+14,  6.21738225e+14, -1.07302720e+14,  5.30843363e+14,\n",
              "        -1.81698043e+15, -7.80751667e+14,  7.43501303e+13, -7.47075311e+14,\n",
              "        -6.55239731e+14,  9.89777546e+14]),\n",
              " array([-3.21488137e+14, -8.61191175e+14,  1.02969926e+15,  2.88151773e+14,\n",
              "         1.63251726e+15,  7.74224022e+13,  1.07668601e+15,  1.08442123e+15,\n",
              "         9.19587762e+14, -1.44725068e+15]),\n",
              " array([-1.46054805e+15,  7.48028908e+14, -8.91907772e+14, -1.80134879e+14,\n",
              "         6.26729257e+14, -8.55134582e+13,  7.72586924e+14,  8.25210858e+14,\n",
              "        -1.77540679e+15,  8.45431723e+14]),\n",
              " array([-1.05762913e+15, -3.33886763e+14, -5.15467649e+14,  1.12107577e+15,\n",
              "         1.11714511e+15,  3.06672637e+14, -8.68464102e+14, -3.76617506e+14,\n",
              "         4.14483453e+14,  9.63114334e+14]),\n",
              " array([ 1.02409260e+15, -2.14474644e+14,  4.25212757e+14, -4.98496092e+13,\n",
              "        -6.12398384e+14, -1.66665996e+14, -7.61665154e+14,  6.60756919e+14,\n",
              "        -8.92047315e+14, -1.18932073e+14]),\n",
              " array([ 7.85273714e+14,  8.22130201e+14, -3.20059068e+13, -3.34367957e+14,\n",
              "         3.09530119e+14,  9.46340001e+14, -1.48385792e+15,  2.28620638e+15,\n",
              "         5.36427867e+14, -1.21156858e+14]),\n",
              " array([ 9.70559947e+14, -4.68046042e+14,  8.47607883e+14,  5.71602135e+14,\n",
              "         3.37017555e+14,  1.08539558e+15,  6.70405727e+14,  7.56167111e+14,\n",
              "        -1.11441156e+15, -6.44725137e+12]),\n",
              " array([-2.67387960e+14, -9.61004909e+14, -1.05957568e+14,  9.09152787e+14,\n",
              "        -1.55697018e+15,  3.96159212e+14,  1.17873349e+15, -1.09640481e+15,\n",
              "        -1.53349523e+14,  4.67532322e+14]),\n",
              " array([ 1.17381703e+15,  4.76347949e+14, -3.37202570e+14,  8.65384370e+14,\n",
              "        -1.29221982e+15,  4.86417975e+13,  7.82247447e+14,  7.28252278e+14,\n",
              "        -2.21516496e+14,  1.15710242e+15]),\n",
              " array([-3.23200856e+14,  1.87352087e+14,  4.54508874e+14,  1.10655835e+15,\n",
              "         5.75483148e+14,  6.67154157e+14,  1.26441485e+15,  4.17522533e+14,\n",
              "        -4.25349502e+14,  1.24521294e+14]),\n",
              " array([-6.24815671e+14,  4.41199110e+14, -1.30768458e+15,  3.46792218e+14,\n",
              "        -8.45040484e+13,  1.32108399e+15, -4.41045350e+14,  4.69175475e+14,\n",
              "         1.02299156e+15, -2.65234449e+14]),\n",
              " array([-1.27857104e+15, -7.41507188e+14,  7.80397855e+12,  4.81044767e+14,\n",
              "         4.07149629e+14, -2.50665242e+14,  6.46869910e+14, -1.61301773e+15,\n",
              "         1.04763926e+15, -5.34121336e+14])]"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "y_preds  # with layer_index=10"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "r_q-ot_gLwuS",
        "outputId": "7b8464d7-7959-415d-e124-84b5bd58b1db"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[array([ 1.01933446e+14,  8.05261820e+13,  1.84602267e+14,  1.05339799e+14,\n",
              "         4.43008251e+13,  7.01698047e+13,  1.09704668e+14, -3.20906777e+13,\n",
              "         2.05806608e+13, -1.88376388e+13]),\n",
              " array([8.56284749e+13, 1.03866768e+14, 5.86992902e+13, 1.19602031e+14,\n",
              "        4.33631128e+13, 7.04106850e+13, 9.68914223e+13, 6.07550786e+13,\n",
              "        4.64704957e+13, 2.06202365e+13]),\n",
              " array([ 8.10123404e+13, -4.36167139e+13,  7.28009635e+13,  1.37008496e+14,\n",
              "        -3.97418926e+13, -3.76108309e+13,  6.91322120e+13,  1.86807056e+13,\n",
              "         9.84340932e+13,  1.13996722e+13]),\n",
              " array([ 7.10702042e+13,  2.73719581e+13,  3.28367932e+13,  9.32909089e+13,\n",
              "         4.87876590e+12,  8.08892370e+13, -4.87724506e+12,  3.46425689e+13,\n",
              "         5.50322772e+13,  4.51992930e+13]),\n",
              " array([ 1.71248540e+13,  7.24101279e+13,  5.77935562e+13,  4.40714346e+13,\n",
              "        -2.38598329e+13, -6.08996426e+12,  2.56322698e+13,  1.47027046e+14,\n",
              "         8.83174080e+13, -4.05753798e+12]),\n",
              " array([ 4.42034166e+13,  1.28657632e+14,  1.20201528e+13,  5.15171595e+13,\n",
              "         1.26924926e+13,  2.97851084e+13,  8.92838520e+13,  7.60187114e+13,\n",
              "        -2.95514310e+13,  1.66563989e+14]),\n",
              " array([ 1.64763199e+14,  1.02381541e+14, -4.58555825e+13, -1.58420066e+13,\n",
              "         8.41023337e+13, -1.44887286e+13,  7.16902926e+13,  3.51148298e+13,\n",
              "         8.51380753e+13,  8.28731283e+13]),\n",
              " array([-1.31687554e+13,  1.23888311e+14,  1.30095847e+12,  8.56100259e+13,\n",
              "         6.02043821e+11,  2.70066695e+13,  8.01992408e+13, -2.71059910e+13,\n",
              "         1.45041360e+14,  4.19042147e+13]),\n",
              " array([-2.78582441e+12,  1.28769774e+14,  9.05923363e+13, -2.12465019e+13,\n",
              "         1.27680663e+14,  7.89139333e+13,  6.84463063e+13, -9.94892586e+12,\n",
              "        -1.97133716e+13,  1.27108983e+14]),\n",
              " array([ 6.56457211e+12, -2.45018981e+13,  1.17633664e+14,  5.44737177e+13,\n",
              "         3.15585542e+13,  7.16579421e+12,  5.10244915e+13,  6.66125383e+13,\n",
              "         6.01845393e+13,  1.14877400e+14]),\n",
              " array([ 7.49667401e+13,  5.73525145e+13,  6.54080644e+12,  1.93291303e+13,\n",
              "         9.84574615e+13, -7.79952246e+12,  1.23451247e+14,  1.41467137e+14,\n",
              "        -2.74515648e+12,  7.17768847e+13]),\n",
              " array([ 2.51578007e+13,  8.00057780e+13,  8.00890434e+13,  3.07631700e+13,\n",
              "         4.18367476e+13,  2.29444873e+13,  9.70663094e+13,  4.95144100e+12,\n",
              "        -1.65312508e+13,  1.14295813e+14]),\n",
              " array([ 7.65069757e+13,  5.62266778e+13, -4.03797177e+13,  8.60321775e+13,\n",
              "         3.06901034e+13,  3.10962513e+13,  1.33936114e+14,  6.20464782e+13,\n",
              "         4.18093224e+13,  5.99277952e+13]),\n",
              " array([ 1.06369739e+14, -3.99179324e+12, -1.51596715e+12, -1.91978222e+13,\n",
              "         3.72978863e+13,  5.71784058e+13,  3.01026595e+13, -5.74090567e+13,\n",
              "        -3.27074081e+13,  1.52518407e+13]),\n",
              " array([ 9.83109294e+13,  1.01260070e+14, -1.47935528e+13,  6.63612155e+13,\n",
              "         6.49341955e+13,  1.61957636e+14,  2.44584400e+11,  5.20817777e+13,\n",
              "        -4.65956863e+13, -8.48852632e+12]),\n",
              " array([ 1.94404117e+13,  6.42582667e+13,  4.84527694e+13,  1.18726095e+14,\n",
              "        -2.29834958e+13,  1.29365275e+14, -4.24836199e+13, -4.39405947e+11,\n",
              "         2.64556283e+13,  8.58806034e+13]),\n",
              " array([ 4.14161925e+13,  1.02015166e+14,  5.62387790e+13,  1.36825903e+14,\n",
              "         5.46611226e+13,  9.44702276e+13,  4.64463604e+13,  2.23044902e+13,\n",
              "         8.29529625e+13, -5.83070188e+13]),\n",
              " array([ 9.73576189e+13, -1.19230766e+12,  1.79080171e+13, -3.28455635e+13,\n",
              "         3.63987793e+13,  7.76766906e+13,  4.07858512e+13, -4.77957736e+13,\n",
              "         5.81324153e+13,  7.81707094e+12]),\n",
              " array([-1.59827045e+10, -5.12527222e+13,  5.83063686e+13,  2.22266478e+13,\n",
              "        -3.33998610e+13,  9.59149671e+13,  2.56517659e+13,  8.80737759e+13,\n",
              "         1.24060548e+14,  8.57061666e+13]),\n",
              " array([ 7.84257504e+13,  7.20967281e+13, -1.22308021e+13,  8.81218150e+12,\n",
              "         7.61607659e+13,  9.31652476e+13,  1.18757522e+14,  5.11901035e+13,\n",
              "         3.45263737e+11,  4.48386870e+13]),\n",
              " array([ 5.67071319e+13, -3.88434718e+13, -3.09872643e+13, -4.37984308e+12,\n",
              "         3.22095383e+12, -6.78212324e+13,  5.18154341e+12,  2.67973197e+13,\n",
              "         1.29102419e+14, -8.70924855e+13]),\n",
              " array([ 5.85705663e+13,  2.50454810e+13, -2.40758423e+13,  2.89141497e+12,\n",
              "         1.01819343e+14,  1.00656954e+14,  5.44579130e+13,  1.22238205e+14,\n",
              "         1.22369947e+13,  7.24027198e+13]),\n",
              " array([ 1.01895007e+14,  1.07798742e+14,  3.44361814e+13, -1.90289071e+13,\n",
              "         4.79454452e+13,  6.72924477e+13,  7.83665490e+13,  2.59107777e+13,\n",
              "         6.55021762e+13,  1.20953210e+13]),\n",
              " array([ 9.79931132e+13,  4.29015821e+13,  5.32411416e+13,  3.34841611e+13,\n",
              "         2.31309677e+13,  6.08242715e+13,  4.37036065e+13,  1.18155800e+13,\n",
              "         2.11821471e+14, -7.64200603e+13]),\n",
              " array([-1.92901693e+13, -3.92834820e+13,  6.90573168e+13, -6.23891820e+13,\n",
              "         1.29694710e+14,  1.30755606e+14, -7.01868047e+13,  4.08050704e+13,\n",
              "         1.19817062e+14,  5.44875240e+13]),\n",
              " array([-8.21677548e+12,  4.78815594e+13,  1.20616169e+14,  3.38274207e+13,\n",
              "         2.94070330e+12,  5.43930819e+13, -1.12264755e+13,  1.26876876e+14,\n",
              "         9.44597170e+13, -8.64083623e+12]),\n",
              " array([ 1.29907326e+13,  6.76754529e+13,  8.08036958e+13,  2.50046294e+13,\n",
              "         3.98848842e+13,  1.09224109e+14,  3.65330795e+13,  5.40565415e+12,\n",
              "         1.93144463e+13, -1.79851203e+13]),\n",
              " array([ 5.66020941e+13,  4.51703115e+13, -6.62467920e+13, -1.33883165e+13,\n",
              "         9.64169217e+13,  1.00087031e+14, -1.80547075e+12, -3.07780297e+13,\n",
              "         9.19251194e+13,  5.14268648e+13]),\n",
              " array([-7.63490825e+12,  6.33718756e+13,  1.12316893e+14,  4.46091902e+13,\n",
              "         8.62714176e+12,  7.25885206e+13, -1.23562781e+13,  1.34885932e+14,\n",
              "         9.35125060e+13,  6.26747271e+13]),\n",
              " array([1.86439076e+13, 1.05406510e+14, 7.81098671e+13, 1.18770367e+14,\n",
              "        9.33130096e+13, 3.62928968e+12, 6.48598105e+13, 3.24233277e+13,\n",
              "        9.10353957e+13, 3.47425237e+13]),\n",
              " array([ 9.24594840e+13,  1.37637550e+14, -4.86599977e+12,  6.54828668e+13,\n",
              "         2.66881598e+12,  3.36676857e+13, -3.43104610e+13, -2.68655245e+13,\n",
              "         9.37697793e+13,  3.87220925e+13]),\n",
              " array([ 5.98203731e+13,  8.95329699e+13,  6.71305127e+13,  6.28423623e+13,\n",
              "        -9.01223891e+13,  3.22620811e+13,  7.12455786e+12,  1.02187960e+14,\n",
              "         2.82136053e+13,  2.95333736e+13]),\n",
              " array([-3.70184755e+13,  4.74921072e+13,  9.45257953e+13, -3.03666663e+13,\n",
              "         5.89590722e+13,  1.26683607e+14,  8.91781344e+13,  7.87453143e+13,\n",
              "         8.37927633e+13, -8.62854329e+13]),\n",
              " array([ 1.46838919e+13,  5.83785219e+13,  6.09624294e+13,  6.15689620e+13,\n",
              "        -4.02929931e+13,  6.19632827e+13,  3.69208287e+12, -2.19382957e+13,\n",
              "        -2.43267045e+13, -3.57937530e+12]),\n",
              " array([ 3.36410477e+13, -1.15804791e+14,  8.32588474e+12,  5.26671186e+13,\n",
              "         6.90874432e+13,  6.44126962e+13,  2.11504433e+14,  4.14272650e+13,\n",
              "         1.49197025e+14,  5.76220881e+13]),\n",
              " array([ 1.65275546e+14, -5.38426720e+13,  5.69238454e+13,  5.18957910e+13,\n",
              "         2.23659064e+13,  1.49364490e+14,  1.90583186e+13,  4.67659750e+13,\n",
              "         1.07280572e+14, -3.43964175e+12]),\n",
              " array([ 7.39482826e+13,  1.43614332e+14,  3.19567648e+13, -5.85578051e+13,\n",
              "         1.18638498e+13,  1.95241391e+13,  9.15856371e+13,  4.08810348e+13,\n",
              "        -6.46385184e+13, -1.92652314e+13]),\n",
              " array([1.35563778e+13, 7.33194076e+13, 5.92554852e+13, 4.49190544e+13,\n",
              "        4.96331715e+13, 4.44740220e+13, 1.59401750e+14, 2.90837326e+13,\n",
              "        3.49237870e+13, 8.16813247e+13]),\n",
              " array([ 5.82924126e+13,  4.21505151e+13,  9.50729788e+13,  2.37189129e+13,\n",
              "        -9.27165005e+12,  1.72628881e+13,  9.42072675e+13,  7.52295038e+13,\n",
              "        -3.07983293e+13,  3.77505286e+13]),\n",
              " array([ 1.42478635e+14,  5.01357242e+12,  3.02666837e+13,  1.35986164e+14,\n",
              "         4.33093366e+13, -1.17721832e+13,  1.04451459e+13,  3.10996962e+13,\n",
              "         8.65650951e+13,  5.52611407e+13]),\n",
              " array([-3.21496425e+13, -1.62034305e+13, -2.12311815e+13, -9.42469907e+12,\n",
              "         1.11620654e+14,  1.60629865e+13,  5.53845656e+13,  4.90544318e+13,\n",
              "         4.21719640e+13,  3.20894747e+13]),\n",
              " array([ 5.60637627e+13,  6.63201127e+13,  5.99550860e+13,  2.35304435e+13,\n",
              "         8.51156401e+13,  8.87850639e+13,  9.00599394e+13,  1.70932054e+14,\n",
              "         2.20361748e+13, -3.35147357e+13]),\n",
              " array([2.10035551e+14, 1.62743809e+14, 3.51032875e+13, 6.66573034e+13,\n",
              "        5.83124329e+12, 1.13410715e+14, 1.99182516e+13, 6.86838659e+13,\n",
              "        3.12020074e+13, 1.09525303e+14]),\n",
              " array([ 6.68970359e+13,  3.86890774e+13, -2.82664162e+13,  2.85749130e+13,\n",
              "         4.86555756e+13,  1.83068937e+13,  5.91016803e+13,  4.41228108e+13,\n",
              "         2.74768409e+13,  7.57619986e+13]),\n",
              " array([ 5.44571314e+13,  1.05226266e+14,  6.23608574e+13, -1.47453875e+14,\n",
              "         1.90099551e+13,  4.84408472e+13,  4.44407255e+13,  2.93516690e+13,\n",
              "         1.36719667e+13,  3.26672687e+13]),\n",
              " array([ 1.41157644e+14,  2.98497156e+13, -1.56152302e+13,  1.08469193e+14,\n",
              "        -2.04951969e+13,  1.40737839e+13,  5.52280607e+13, -3.58945184e+13,\n",
              "         7.81645319e+13,  6.62592410e+13]),\n",
              " array([ 2.27312162e+14,  8.29342983e+13, -3.12117239e+13,  4.24578812e+13,\n",
              "         1.56259142e+13,  5.17078570e+13,  8.89899011e+12, -3.57794587e+12,\n",
              "         1.07743381e+13,  9.03697949e+13]),\n",
              " array([-7.92191072e+12,  7.52048121e+13,  7.37405495e+13,  7.90865812e+13,\n",
              "         4.33324899e+13,  1.05175673e+14,  3.59744137e+13,  1.34444827e+14,\n",
              "        -3.10854854e+13,  4.76695953e+13]),\n",
              " array([-7.15381405e+13,  4.79638355e+13,  8.35388443e+13,  7.91248039e+13,\n",
              "         1.09625251e+14,  1.27775919e+13,  5.49659839e+13, -9.00466885e+12,\n",
              "         1.66074769e+13, -1.11861006e+13]),\n",
              " array([-3.73916828e+13,  1.49412445e+14, -1.36174433e+12,  4.34561658e+13,\n",
              "        -9.21512206e+13,  1.18419395e+14,  1.83234860e+14,  6.08538434e+13,\n",
              "         1.30660479e+12,  1.16368239e+12]),\n",
              " array([ 1.09177441e+14,  1.60483727e+14,  9.13749121e+13, -3.94192811e+13,\n",
              "         5.98011771e+13,  3.59079557e+13,  2.71668845e+13,  6.99774980e+12,\n",
              "         1.04698681e+14,  4.49301140e+13]),\n",
              " array([ 6.85660621e+13, -3.41335580e+13, -1.56373156e+13,  3.01719371e+13,\n",
              "         5.58122032e+13,  3.82322936e+13, -2.53948594e+13, -8.59234944e+12,\n",
              "         5.49011318e+13, -1.80437511e+13]),\n",
              " array([-1.09567461e+13,  3.54512907e+13, -3.49068979e+13,  8.53824443e+13,\n",
              "         6.33745231e+13, -1.07288022e+14,  2.04532220e+13,  7.24230595e+13,\n",
              "         1.46362449e+14,  1.01628210e+14]),\n",
              " array([-6.54936422e+12,  1.28340341e+14, -1.84465039e+13,  1.77330564e+14,\n",
              "         1.90043207e+14,  7.03426251e+13, -5.71095123e+13,  1.04209107e+14,\n",
              "         6.44143939e+13,  7.22702434e+13]),\n",
              " array([1.35586282e+14, 1.72696379e+13, 4.63836299e+13, 2.37593563e+13,\n",
              "        1.87026053e+14, 1.86082736e+14, 6.48058118e+13, 8.38147860e+13,\n",
              "        1.03785082e+13, 1.20116090e+14]),\n",
              " array([-3.25260179e+13,  1.69466556e+13,  1.31132325e+13,  8.66687872e+12,\n",
              "         5.74949981e+13,  6.30123537e+13, -3.03695498e+13,  6.64052068e+13,\n",
              "         5.28333540e+13,  4.65264836e+13]),\n",
              " array([5.84925380e+13, 2.32146501e+13, 6.02752811e+13, 1.36117262e+14,\n",
              "        4.42768735e+13, 1.46961927e+14, 1.12264226e+13, 7.57210544e+13,\n",
              "        1.13378241e+14, 4.76865284e+13]),\n",
              " array([ 6.21194762e+13,  6.38596152e+13,  2.80143626e+13,  6.14777805e+13,\n",
              "         6.45001225e+13,  2.43223622e+13,  8.97591766e+13, -1.89209174e+12,\n",
              "         6.39676159e+13,  6.67416764e+13]),\n",
              " array([7.20355960e+13, 8.60003676e+13, 4.63487038e+13, 3.95418414e+13,\n",
              "        8.97248860e+13, 9.79049370e+13, 1.88518006e+13, 1.03769261e+13,\n",
              "        5.69147173e+13, 5.06955075e+13]),\n",
              " array([ 5.45905629e+13,  9.04314349e+13, -4.54435000e+13,  9.42113596e+13,\n",
              "         9.48716698e+13, -1.03480877e+13,  3.72438911e+13,  7.63985249e+13,\n",
              "         5.10801829e+13,  3.12761870e+13]),\n",
              " array([ 2.03925423e+13,  6.99413391e+13,  8.52718993e+13,  8.23635220e+12,\n",
              "         5.10507555e+13, -1.26098787e+14, -2.87664853e+13,  1.27346217e+14,\n",
              "         4.82758922e+13,  7.60949250e+13]),\n",
              " array([ 7.55106753e+13,  4.07584181e+13,  1.89346407e+13,  8.19165889e+13,\n",
              "         5.98150741e+13, -5.49798787e+12, -6.74954869e+13, -3.81863023e+12,\n",
              "        -5.02917404e+11,  9.19690651e+13]),\n",
              " array([ 7.68645321e+13,  8.04695595e+13,  1.89500350e+13,  8.14709429e+12,\n",
              "         5.67348793e+13,  4.19349705e+13,  7.61589979e+13,  4.43628338e+13,\n",
              "         7.52115250e+13, -2.07555786e+13]),\n",
              " array([-4.65694686e+13,  6.01046947e+13,  3.34366911e+13,  5.77597531e+13,\n",
              "         1.00119611e+12,  6.63496131e+13,  4.54893159e+13,  5.89729604e+13,\n",
              "         6.64416149e+13,  1.07842485e+14]),\n",
              " array([ 5.22037925e+13,  4.04488985e+13, -6.58964344e+13,  6.17422274e+13,\n",
              "         1.00233143e+14,  8.42349285e+13,  8.65658544e+12,  4.97179184e+13,\n",
              "         6.43798778e+13,  1.79785331e+13]),\n",
              " array([ 1.07318792e+14, -3.04867506e+13,  2.84339437e+13,  2.04845792e+13,\n",
              "         5.23493695e+13, -6.46396916e+12,  4.70748011e+13,  1.97608687e+12,\n",
              "         1.11941670e+14,  1.03870291e+14]),\n",
              " array([3.02632913e+13, 1.00078408e+14, 2.38567335e+13, 2.27084866e+13,\n",
              "        7.17973643e+13, 2.44920238e+14, 4.58948183e+13, 7.07445838e+13,\n",
              "        2.24678517e+13, 5.09293040e+13]),\n",
              " array([-1.52498863e+12,  3.98306281e+13, -1.08850087e+12,  4.18621895e+12,\n",
              "         1.02479331e+14,  5.58947151e+13, -6.50939698e+13,  5.63727759e+13,\n",
              "        -4.15892931e+13,  3.60291267e+13]),\n",
              " array([-9.39087772e+11,  8.77806079e+13,  1.30855331e+14, -1.72439567e+14,\n",
              "         3.82980642e+13,  8.81552670e+13,  4.53719256e+13, -8.23476645e+12,\n",
              "         3.81609210e+13,  2.64803470e+13]),\n",
              " array([ 7.75178240e+13,  1.84363949e+13,  9.16993422e+12,  1.01099581e+13,\n",
              "         7.96830713e+13,  3.70592342e+13, -3.08693782e+13,  1.73567082e+14,\n",
              "         7.31507729e+13,  5.77763061e+13]),\n",
              " array([8.80268078e+12, 6.57329766e+13, 7.04772416e+13, 3.64744854e+13,\n",
              "        1.58777650e+14, 4.77782604e+13, 8.85767260e+13, 8.83823727e+13,\n",
              "        4.73127413e+13, 1.15821309e+14]),\n",
              " array([1.09021499e+14, 2.38341982e+13, 8.88549030e+13, 4.03958669e+13,\n",
              "        1.93575730e+13, 4.87437870e+13, 6.44142525e+13, 1.36874762e+14,\n",
              "        8.12173388e+13, 1.44421208e+12]),\n",
              " array([8.74263324e+13, 1.35821820e+14, 3.47594033e+13, 1.26766524e+13,\n",
              "        3.18682041e+13, 8.07457072e+13, 2.81369682e+13, 8.40490279e+13,\n",
              "        3.42758227e+13, 1.00789191e+14]),\n",
              " array([ 5.64468893e+13,  6.92335799e+13, -5.09831447e+12,  1.92230534e+13,\n",
              "        -5.78805036e+13,  4.63281345e+13,  4.27424052e+13,  9.21562643e+13,\n",
              "         3.81389901e+13,  1.69470365e+14]),\n",
              " array([ 1.14788027e+14,  1.80074367e+14,  2.15738607e+13,  3.14013876e+13,\n",
              "         3.06532057e+11,  8.21182986e+13,  3.93552553e+13,  7.42698939e+13,\n",
              "         1.27212878e+14, -2.78200726e+13]),\n",
              " array([ 9.40652922e+13,  6.83462059e+13, -7.02702287e+12,  1.68920999e+13,\n",
              "         6.10439013e+13,  4.78481526e+13,  1.41714922e+14,  5.28855338e+13,\n",
              "         7.38581037e+13,  3.92468726e+13]),\n",
              " array([-2.36353219e+13,  6.82835322e+13,  3.22699327e+12, -4.07126979e+12,\n",
              "         2.54955272e+13,  2.37061976e+13,  4.65012276e+13,  8.62265978e+13,\n",
              "        -1.04387028e+14,  1.18461537e+14]),\n",
              " array([ 1.34987549e+13,  9.66292419e+13,  8.63843965e+13,  1.38844213e+13,\n",
              "        -3.47830195e+13, -9.70591845e+12,  5.14168541e+13, -8.15682825e+12,\n",
              "        -2.39688667e+13,  3.96584030e+13]),\n",
              " array([-1.22417823e+13,  4.20875353e+13,  8.44544095e+13,  4.44988433e+13,\n",
              "         7.01656924e+13,  2.19667141e+13, -1.39463144e+13,  3.83866312e+13,\n",
              "        -4.75266743e+13,  1.38120850e+14]),\n",
              " array([ 5.22004598e+13,  4.42831736e+13,  3.73873147e+13,  9.00619824e+13,\n",
              "         1.96145380e+13,  1.48175417e+14,  5.12272282e+13,  9.07105797e+13,\n",
              "         4.85067864e+13, -4.84860876e+13]),\n",
              " array([ 3.82726927e+13,  7.06696564e+13, -2.42314505e+13,  1.40656243e+13,\n",
              "        -3.21226950e+13,  8.60947571e+13,  3.45126725e+13,  4.05415661e+13,\n",
              "         3.89415615e+13, -9.30848332e+12]),\n",
              " array([2.49732744e+14, 3.66226415e+13, 9.79862661e+13, 9.37816199e+13,\n",
              "        1.69575114e+13, 9.98188259e+13, 1.80283139e+13, 7.36676053e+13,\n",
              "        1.27292852e+14, 1.45362141e+14]),\n",
              " array([ 4.63841255e+13,  6.53373318e+13,  8.33951674e+13, -2.24672820e+13,\n",
              "        -3.76750042e+13,  5.93712488e+13,  8.66292550e+13,  1.45729265e+13,\n",
              "         6.15203785e+13,  5.99032489e+13]),\n",
              " array([ 3.52731545e+13, -1.71826791e+13,  1.32127338e+13,  5.17280794e+13,\n",
              "         4.88171970e+12,  5.02420427e+13, -8.80421177e+13, -7.49968179e+12,\n",
              "         3.02085745e+13,  3.85991453e+13]),\n",
              " array([-3.96309042e+13, -5.07473678e+12,  8.54845863e+13,  6.40847512e+13,\n",
              "         7.40097716e+13,  7.41434305e+13,  5.52483851e+13,  1.13966882e+14,\n",
              "         1.05933840e+14,  5.16294325e+12]),\n",
              " array([ 5.09088219e+13,  3.55513473e+13,  6.69150383e+13,  5.90834700e+13,\n",
              "         1.54004689e+14,  5.46710670e+13, -9.80449692e+12,  4.41524414e+13,\n",
              "         3.86477790e+13,  1.40180321e+14]),\n",
              " array([ 2.51868370e+13, -1.18096499e+13,  1.24440421e+13,  9.10759830e+13,\n",
              "         1.24619062e+14,  3.31801506e+13, -4.18464808e+13,  4.17953993e+13,\n",
              "         3.12912830e+13, -3.36913925e+13]),\n",
              " array([ 7.55116334e+13,  2.19217358e+13,  4.62371409e+13,  6.95962017e+13,\n",
              "         2.01874826e+13, -2.17485664e+13, -1.51647892e+13,  7.38466798e+13,\n",
              "         1.16408359e+14,  9.02044484e+13]),\n",
              " array([9.10310976e+13, 1.10524273e+13, 2.50073008e+13, 8.73498051e+13,\n",
              "        1.77033016e+13, 6.53558355e+13, 4.49764231e+13, 7.94354569e+13,\n",
              "        5.25050666e+13, 5.60258630e+13]),\n",
              " array([ 4.38220959e+12,  3.36687184e+13,  6.96918585e+13,  6.64352018e+13,\n",
              "        -1.68322498e+13, -2.07996740e+13,  8.93022303e+13,  4.10641062e+13,\n",
              "         1.62291228e+13,  2.15691580e+13])]"
            ]
          },
          "metadata": {},
          "execution_count": 27
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "y_preds  # with layer_index=16"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hNwZSgkFV0jz",
        "outputId": "23a1ba6a-90e8-4065-b97a-a8e9749b3638"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[array([-4.45118357e+11,  9.53812006e+11,  1.14939835e+13, -2.93997318e+12,\n",
              "        -1.47145224e+13, -2.38411666e+12, -1.14178486e+13,  5.05004536e+12,\n",
              "         2.26684350e+12,  1.19420816e+13]),\n",
              " array([ 1.62113400e+13, -5.25175842e+12, -1.31925641e+13, -9.29885629e+12,\n",
              "        -1.70820983e+13, -1.71609928e+13,  5.67826318e+12, -8.98114894e+11,\n",
              "         2.39841036e+13, -8.73528483e+11]),\n",
              " array([-1.55433208e+13,  1.50842727e+13,  1.59793258e+13,  1.96340060e+13,\n",
              "         3.51092186e+12, -2.57553615e+13, -6.75088408e+12, -2.89751580e+12,\n",
              "        -1.77541530e+11, -2.78513684e+13]),\n",
              " array([-1.33739206e+13, -6.58147150e+12, -1.75108409e+13,  8.88851986e+12,\n",
              "        -4.30321150e+12,  1.26082120e+13, -1.26120838e+13,  2.99891841e+13,\n",
              "        -1.54913559e+13, -7.10903588e+12]),\n",
              " array([ 2.31293479e+13, -2.05318409e+13, -3.33764721e+12, -2.56637258e+13,\n",
              "        -1.29731125e+13,  2.03272038e+12, -3.92356638e+12, -2.28697099e+13,\n",
              "        -3.87492551e+13, -8.11679912e+12]),\n",
              " array([-8.41310168e+12,  4.76237454e+12,  3.04915600e+12, -2.05724250e+13,\n",
              "        -1.16775452e+13,  2.46941576e+13,  8.79558307e+12, -1.71076682e+13,\n",
              "        -2.91238533e+13,  2.06175663e+11]),\n",
              " array([ 1.92195578e+12,  2.22628669e+11, -4.69569662e+12,  8.54496107e+12,\n",
              "         1.58223334e+13,  1.91413559e+11,  6.01323521e+12, -4.09402113e+12,\n",
              "        -1.65682908e+13, -9.52359501e+12]),\n",
              " array([ 2.62846671e+13,  1.71607079e+13,  1.10079734e+13, -9.04729671e+12,\n",
              "        -7.16248794e+12, -1.41019714e+13,  1.90483889e+12,  5.91560521e+12,\n",
              "         6.46094863e+12, -2.73376833e+13]),\n",
              " array([-2.23810900e+13, -1.11301106e+13, -1.47741716e+13,  1.70521630e+12,\n",
              "        -3.50103033e+12,  1.24491528e+12, -6.28692945e+12, -2.47196666e+13,\n",
              "         4.84606574e+12, -3.31359020e+13]),\n",
              " array([-2.20351158e+13, -1.09141151e+13, -4.07552615e+12,  5.61994669e+12,\n",
              "         1.22100859e+13,  2.09630409e+12, -1.46538466e+13, -1.01199630e+13,\n",
              "        -1.32244359e+13,  1.85690414e+12]),\n",
              " array([ 3.60096708e+12,  2.03281732e+13, -8.84999311e+12,  8.19370712e+12,\n",
              "         3.08472891e+12, -2.17356486e+13, -2.70397622e+12, -2.42549305e+13,\n",
              "        -1.20527135e+13,  3.00965899e+10]),\n",
              " array([-1.60108475e+13,  1.29849875e+13, -9.78516302e+12, -2.56036434e+13,\n",
              "        -2.58604396e+13, -1.25661527e+13, -1.76045574e+13, -6.60587615e+12,\n",
              "        -1.70698624e+13, -2.27702349e+13]),\n",
              " array([-4.22969510e+13,  2.08606619e+13, -7.83577647e+12, -1.58706601e+13,\n",
              "         2.37403670e+12, -1.01208030e+13,  3.25455963e+11,  2.63073275e+12,\n",
              "         3.93002037e+10, -1.88596618e+12]),\n",
              " array([ 1.44170995e+12,  1.12932558e+13,  1.34156736e+13, -9.38229327e+12,\n",
              "        -1.78231918e+13, -1.27647896e+13, -1.95952431e+12, -6.40080302e+12,\n",
              "        -1.16399881e+13, -2.80984124e+13]),\n",
              " array([-3.71418049e+13, -2.51894899e+13,  9.47023915e+12, -4.01527016e+12,\n",
              "        -6.99524596e+12,  9.90706014e+12, -1.65623131e+13, -1.33764569e+12,\n",
              "        -1.57534290e+13,  1.90312221e+12]),\n",
              " array([ 3.03369232e+12,  1.67475910e+13, -8.72228117e+11, -7.02473729e+12,\n",
              "        -8.49247964e+12, -1.62228171e+12, -8.69548145e+12,  1.95243970e+13,\n",
              "         1.05696671e+13, -5.33667492e+12]),\n",
              " array([ 7.26429481e+12, -2.02735993e+13, -8.65604341e+12,  4.79995756e+12,\n",
              "        -7.38573102e+12,  1.07595528e+13,  8.52543045e+12, -1.91723879e+13,\n",
              "        -2.70751937e+13,  7.94506847e+12]),\n",
              " array([ 1.19642838e+13,  3.39517807e+12, -4.31803324e+12, -4.29467982e+12,\n",
              "         1.16759602e+13, -1.51200397e+13, -1.32123601e+13, -1.89677898e+13,\n",
              "         1.36309575e+13, -1.34857068e+13]),\n",
              " array([ 3.94486263e+12,  2.01453163e+13,  3.40256279e+12,  3.22246540e+11,\n",
              "        -1.83298042e+13, -2.22018617e+13,  8.47000111e+12,  5.42878417e+12,\n",
              "        -1.57223200e+13,  9.85630653e+12]),\n",
              " array([-1.06879567e+13,  1.68837111e+13,  8.86176078e+11, -1.74179212e+13,\n",
              "         2.00114792e+11, -1.14866997e+13, -1.84179349e+13, -4.61302109e+12,\n",
              "        -9.64146916e+12, -1.42008134e+13]),\n",
              " array([-6.40758821e+12, -1.87152238e+13, -1.30131323e+13, -3.56002422e+12,\n",
              "         5.68471270e+12, -3.75744198e+12, -2.28382785e+13, -5.11603591e+12,\n",
              "        -2.89290811e+13, -1.41111145e+13]),\n",
              " array([ 5.32948457e+11, -3.49662009e+13, -1.08004397e+13, -5.36883858e+12,\n",
              "         3.15235169e+12,  7.49016847e+12, -2.30215054e+13, -5.31976181e+11,\n",
              "         1.00581568e+13, -1.48210476e+13]),\n",
              " array([ 3.53854841e+12, -2.10609363e+13, -1.24601287e+13, -2.58727140e+13,\n",
              "         1.56725095e+12,  3.40707557e+13,  1.79748591e+12, -4.46011188e+12,\n",
              "         3.25701694e+12, -1.84470453e+13]),\n",
              " array([-2.52966990e+13, -6.75979506e+12, -3.38756157e+13, -2.36818704e+13,\n",
              "        -3.10486173e+13, -1.53753702e+13, -2.73011732e+13, -1.04226098e+12,\n",
              "        -2.43907568e+13, -1.08959214e+13]),\n",
              " array([-6.51193166e+12,  5.82058033e+12, -2.24043341e+13,  1.22603922e+13,\n",
              "        -1.47148062e+13, -1.27315831e+13, -1.21893408e+13, -1.31944554e+13,\n",
              "        -1.77088055e+12, -9.44913043e+12]),\n",
              " array([-1.27140503e+13, -3.71532273e+13, -9.29699886e+11, -2.97341272e+13,\n",
              "        -3.26991871e+12, -2.80933647e+13, -3.21855820e+12, -9.48115493e+12,\n",
              "        -2.53490495e+13, -1.12126288e+13]),\n",
              " array([-1.16793571e+13, -2.85207638e+13,  1.14194073e+13, -4.65580157e+13,\n",
              "        -8.16513890e+12, -4.85092265e+12, -3.08861309e+12, -3.78915155e+12,\n",
              "        -1.91623261e+13,  5.33065807e+12]),\n",
              " array([ 6.82738217e+12, -9.45915296e+12,  5.75102169e+12, -7.61848871e+12,\n",
              "        -1.84434872e+13, -3.60095340e+12,  3.62940582e+12,  2.63110729e+12,\n",
              "        -3.12391559e+13,  5.74126455e+12]),\n",
              " array([ 7.48603214e+12,  1.21045224e+13,  1.04307496e+13, -5.17211159e+12,\n",
              "        -3.16632676e+12,  1.01398423e+13, -1.07832421e+13, -1.50550295e+13,\n",
              "        -1.41925572e+13, -1.71110083e+13]),\n",
              " array([ 2.96148680e+12, -4.04516630e+12, -1.20195869e+13,  1.19575027e+12,\n",
              "        -4.47232617e+12,  1.16014487e+13,  1.23071470e+13, -9.79248301e+12,\n",
              "         5.56975808e+12, -2.18625290e+13]),\n",
              " array([-2.05222848e+13,  6.43196041e+12, -1.59235690e+13, -4.66777293e+13,\n",
              "        -1.56102860e+13, -2.10329386e+13, -2.17635102e+13, -1.13817407e+13,\n",
              "        -2.63922059e+13, -1.22174490e+13]),\n",
              " array([ 4.71212785e+12, -1.67963950e+13,  1.15412456e+13, -1.97637310e+13,\n",
              "        -1.32522312e+13, -6.15508948e+12, -2.52800290e+13,  1.16616809e+13,\n",
              "        -1.08221896e+13, -6.08994894e+12]),\n",
              " array([-1.55942007e+13, -1.57293763e+13,  5.11062773e+12, -2.35622664e+13,\n",
              "        -5.50483445e+12,  1.30826036e+13,  7.93303947e+12, -1.14657251e+13,\n",
              "        -1.06541919e+13, -1.70332617e+13]),\n",
              " array([ 1.69074613e+13, -1.96323246e+12,  2.27235230e+12, -7.06330152e+12,\n",
              "        -6.94409714e+12, -1.15014186e+13,  9.69286732e+12,  3.11857657e+12,\n",
              "        -1.06138258e+13, -2.87843086e+13]),\n",
              " array([ 4.67780966e+12,  7.42867718e+11,  6.27907393e+12, -4.03927465e+13,\n",
              "        -1.14322697e+13,  1.04892997e+13,  2.98044299e+12, -3.11000629e+12,\n",
              "        -7.23932361e+12,  1.60212655e+13]),\n",
              " array([-2.49406634e+13, -7.32899316e+12,  1.31621837e+13, -2.61967397e+13,\n",
              "        -3.41585589e+13, -2.71690928e+12,  2.65190606e+13,  2.75094233e+12,\n",
              "        -8.93638597e+12, -3.33162026e+13]),\n",
              " array([-3.97164445e+12,  4.46855094e+12, -1.05505672e+13, -9.91730032e+11,\n",
              "         1.80023569e+13, -1.81505651e+13, -2.54266079e+13, -2.14754355e+13,\n",
              "         1.72935014e+12, -1.69785133e+12]),\n",
              " array([-1.04132035e+13, -1.06239616e+13, -7.39207269e+12, -4.38890267e+12,\n",
              "        -1.13844101e+13, -1.65270389e+13, -4.04525168e+13, -1.55874271e+13,\n",
              "         5.38826156e+12,  7.86463502e+12]),\n",
              " array([ 5.46794281e+11, -2.15858500e+13,  1.37565088e+13,  1.59391875e+13,\n",
              "        -1.24702764e+13, -2.31365271e+13,  5.32246157e+12, -2.90811458e+13,\n",
              "        -8.07868059e+12, -2.06921171e+13]),\n",
              " array([ 5.24606053e+12, -1.20833092e+12, -9.20587677e+12, -7.92047151e+12,\n",
              "         1.02016427e+13,  5.00473018e+12,  9.40828180e+12,  7.52404074e+12,\n",
              "        -1.58215751e+13, -3.08159862e+13]),\n",
              " array([-2.44103093e+12, -8.20807833e+12, -1.95346683e+12, -9.76552363e+11,\n",
              "        -1.35627341e+12, -8.39124627e+12, -2.61061155e+13,  1.88201498e+13,\n",
              "         1.66733976e+13, -3.62279769e+13]),\n",
              " array([-3.10942619e+13, -1.37430992e+13, -1.94613858e+12, -2.30367300e+13,\n",
              "        -1.57549988e+13, -1.60307991e+13, -1.24832252e+13, -1.22108678e+13,\n",
              "         1.37311464e+12, -6.34512765e+12]),\n",
              " array([-4.54507618e+12,  1.25594175e+13,  1.36901100e+13,  9.80482241e+12,\n",
              "        -1.35524884e+13, -3.71724027e+12, -2.28881173e+13, -1.11954198e+13,\n",
              "        -9.68068724e+12, -1.98923609e+13]),\n",
              " array([-1.42734129e+13, -1.03787539e+13,  1.29816216e+13,  2.38232669e+13,\n",
              "         1.50182374e+13, -1.87515730e+13,  1.28450711e+13, -1.58094709e+13,\n",
              "        -8.50943104e+11, -1.69025579e+13]),\n",
              " array([-1.71760111e+13, -6.49054256e+12, -9.49251296e+12, -1.99245138e+13,\n",
              "        -1.93774497e+13, -2.37130972e+12, -1.53515787e+13,  1.56418543e+13,\n",
              "        -1.06220527e+13,  5.76747584e+12]),\n",
              " array([-1.16295502e+12, -1.02136657e+13,  5.35401938e+12,  9.82879630e+12,\n",
              "        -1.27180089e+11, -1.25955866e+12, -1.18648603e+13, -5.28709204e+12,\n",
              "         1.48913698e+13, -4.38414500e+12]),\n",
              " array([-7.95836433e+12, -3.15473567e+11,  1.19049682e+11, -2.83234117e+13,\n",
              "        -1.43955750e+12,  2.74566327e+13, -4.39428586e+13,  1.35103466e+13,\n",
              "        -1.25006648e+12, -8.76277426e+12]),\n",
              " array([ 1.20633024e+13, -3.33735459e+13, -5.20027740e+12, -6.41546478e+12,\n",
              "         6.31445288e+11, -1.79055444e+13,  4.44221177e+12,  7.62521468e+12,\n",
              "        -1.34279210e+13,  2.89649133e+10]),\n",
              " array([-1.30713480e+13,  3.59723944e+12, -2.73570748e+13, -1.98756304e+13,\n",
              "         1.47182532e+13,  1.27597578e+13,  8.60027761e+12, -2.08289437e+12,\n",
              "         7.76509226e+12,  1.42652123e+13]),\n",
              " array([-1.17765174e+13, -1.65463542e+13,  7.73018724e+12,  8.32165067e+11,\n",
              "        -1.86574157e+13, -2.66437537e+12,  3.83485261e+12,  5.52202153e+12,\n",
              "         2.85756750e+12, -1.02643131e+13]),\n",
              " array([-1.97289429e+13, -8.57803870e+11, -2.51851181e+13, -1.08102345e+13,\n",
              "        -1.18010981e+13, -1.18076967e+13, -1.80956836e+12, -4.73564768e+12,\n",
              "        -8.08443901e+12,  5.19101517e+12]),\n",
              " array([-6.14829486e+11,  1.58509574e+13, -2.49665772e+13, -8.12954164e+11,\n",
              "        -1.31961387e+12, -1.30239075e+13, -2.47363409e+13, -1.89180833e+13,\n",
              "        -1.80066087e+13, -6.99975183e+12]),\n",
              " array([ 1.61403094e+13, -8.03736292e+12, -1.37639445e+13, -3.59935890e+13,\n",
              "        -8.91399233e+12, -4.50056264e+12, -2.75681526e+12, -2.28704252e+13,\n",
              "         1.04626944e+10, -1.28205814e+13]),\n",
              " array([ 1.70071683e+12, -8.13825948e+12, -2.98710685e+13,  5.86144479e+12,\n",
              "        -2.22133009e+13, -2.68313237e+13, -7.77365128e+12,  2.18676463e+12,\n",
              "        -1.32696717e+13, -2.86914886e+13]),\n",
              " array([ 2.47504604e+12, -1.87813221e+13, -9.06501298e+12,  9.02143039e+12,\n",
              "        -2.20015434e+13, -3.24967876e+12, -8.63909396e+12, -2.97565062e+12,\n",
              "         1.64347171e+13, -4.12820564e+12]),\n",
              " array([ 1.28717799e+13, -3.55425308e+12, -5.80426900e+12, -7.25027957e+12,\n",
              "         1.12876349e+13, -4.25633631e+12, -9.04763179e+12, -1.77890010e+13,\n",
              "        -1.85132444e+13, -2.05638871e+12]),\n",
              " array([ 3.40647042e+12, -4.21243366e+11,  2.83250044e+12, -1.41567818e+13,\n",
              "        -1.78875542e+12,  1.05790600e+13,  2.88628716e+12, -1.40097898e+12,\n",
              "         1.23588002e+13, -1.94326491e+13]),\n",
              " array([-1.43777696e+13, -1.01525170e+13,  8.99747782e+12, -6.20064369e+12,\n",
              "        -1.15217513e+13, -3.00081310e+13,  1.82275905e+13,  5.10356253e+12,\n",
              "        -1.84870295e+13,  2.17779868e+12]),\n",
              " array([-1.32590660e+13,  1.07740117e+13, -2.18411341e+13, -2.86286128e+13,\n",
              "        -1.86206097e+13, -2.79121275e+13, -2.19906851e+11,  6.23914558e+12,\n",
              "         1.04223295e+13, -1.60682204e+13]),\n",
              " array([ 2.09656809e+11, -1.78345531e+12, -1.43971694e+13, -1.99600684e+13,\n",
              "        -6.49670868e+12,  3.28977752e+12,  2.03573651e+12, -1.14860702e+13,\n",
              "        -3.17806180e+12, -1.93813274e+13]),\n",
              " array([-3.56734008e+13,  2.54774039e+13,  1.16097090e+13,  1.97979381e+13,\n",
              "        -2.65096196e+13, -1.69991055e+13, -3.44029441e+13, -5.40949078e+11,\n",
              "        -1.73093852e+13,  2.72468333e+12]),\n",
              " array([-7.74739715e+12,  4.73224344e+12, -1.83891095e+13, -2.78027958e+12,\n",
              "        -2.93124368e+13, -2.39667733e+13, -3.96787004e+12, -5.11929133e+12,\n",
              "        -1.00654678e+13, -4.55572304e+12]),\n",
              " array([-2.22882417e+13,  1.95004223e+13, -3.27996371e+11,  2.05470623e+13,\n",
              "         9.37123100e+12, -1.36425466e+13,  1.00397412e+13, -3.53773658e+12,\n",
              "        -1.57332628e+12,  1.63739263e+13]),\n",
              " array([-5.44559000e+12, -1.85203431e+13,  1.04361353e+13, -2.75343998e+11,\n",
              "         4.47844141e+11, -1.38443304e+13,  5.45383968e+12,  1.21767623e+13,\n",
              "        -1.33389823e+13, -1.46029557e+13]),\n",
              " array([-1.44884441e+13, -3.69634609e+12, -7.76096511e+12,  4.41482591e+11,\n",
              "        -1.30139532e+13, -1.29491073e+13, -6.15514720e+11,  1.07546451e+13,\n",
              "        -1.15512535e+13, -1.70065664e+13]),\n",
              " array([ 5.47352162e+12,  9.21538142e+12, -2.76517309e+13,  8.27413690e+12,\n",
              "        -3.42849455e+12, -2.97828669e+12, -1.47189683e+13, -1.91764430e+13,\n",
              "        -2.74892683e+13, -1.15722926e+13]),\n",
              " array([-4.66240658e+12, -1.50292893e+13, -1.00055902e+13,  2.76107688e+12,\n",
              "         1.10750494e+13,  1.30474244e+13,  6.76409110e+12,  7.67548266e+12,\n",
              "         1.53291040e+13, -1.01715378e+13]),\n",
              " array([-5.53993571e+12,  1.10940907e+13,  1.50818150e+13, -5.04097135e+12,\n",
              "        -2.89914114e+13, -1.71337614e+13, -9.66164273e+12, -1.25848121e+13,\n",
              "         1.43076073e+13, -1.79490109e+13]),\n",
              " array([-2.58422110e+13, -3.11950504e+12,  9.16038220e+12, -1.07369711e+13,\n",
              "         6.47556458e+12, -2.28987913e+13, -1.27501958e+12,  2.97107071e+11,\n",
              "        -4.49090806e+12, -2.36758749e+13]),\n",
              " array([-1.27262065e+13, -2.10236110e+11, -7.46552502e+12, -1.61013401e+13,\n",
              "        -1.71485658e+13,  3.06234063e+12, -1.52706855e+13, -2.01088153e+13,\n",
              "         6.78786954e+11, -4.36353068e+12]),\n",
              " array([ 8.44295736e+12,  3.27364555e+12, -2.57576091e+12,  3.09607656e+13,\n",
              "         1.37196555e+13, -2.80306626e+13, -1.27805705e+13, -2.84337356e+12,\n",
              "        -1.09716767e+13,  8.68264606e+12]),\n",
              " array([-4.60621468e+12, -1.71685948e+13, -1.98751613e+13,  8.63392816e+12,\n",
              "        -1.94206151e+13,  2.03637446e+13, -1.03890772e+13, -2.92721369e+11,\n",
              "         6.56705376e+12, -1.15205721e+12]),\n",
              " array([-1.44977032e+13,  9.30297720e+11,  5.05414476e+12, -1.70314920e+13,\n",
              "        -2.50921886e+13,  1.76002634e+12,  1.09255061e+12,  4.02559574e+12,\n",
              "        -9.44794946e+12, -3.14273870e+13]),\n",
              " array([-3.80547053e+12,  4.21198295e+12, -1.02208331e+12,  1.19181956e+13,\n",
              "         4.95596768e+12,  2.78989324e+11,  6.88263820e+12, -8.51044017e+12,\n",
              "        -1.82302964e+13, -1.65132425e+13]),\n",
              " array([-5.99370091e+12,  1.19997404e+13, -1.19755651e+13, -1.13577496e+13,\n",
              "        -1.13106236e+13,  1.60415933e+13,  2.35854439e+13, -6.55787438e+12,\n",
              "        -6.85699887e+12, -1.23355062e+13]),\n",
              " array([-1.89768295e+13, -8.16632967e+12,  4.63447998e+11, -8.60657089e+12,\n",
              "        -1.25420109e+13,  3.00265754e+12, -2.44252463e+13, -2.96564602e+12,\n",
              "        -1.70820669e+13, -2.17313367e+13]),\n",
              " array([ 1.22871023e+12, -1.70399016e+13, -2.06152048e+13,  3.05015497e+12,\n",
              "        -1.36909704e+13, -7.14906311e+12,  6.96314913e+12, -2.12889253e+13,\n",
              "        -2.68946025e+12,  4.05697526e+12]),\n",
              " array([-6.32643522e+12, -2.11804121e+13, -2.16336029e+13,  2.51603745e+13,\n",
              "         6.12431302e+12, -2.05667523e+13, -1.33602032e+13,  1.10955647e+13,\n",
              "        -4.43057249e+12, -1.86294192e+13]),\n",
              " array([-3.15542760e+11,  3.25718660e+12, -8.32941355e+12, -1.52220091e+13,\n",
              "        -7.75682456e+12, -1.12870651e+13,  2.57221585e+13, -3.81724082e+12,\n",
              "        -3.28179985e+13, -1.10835360e+13]),\n",
              " array([ 3.11608714e+12, -3.86882169e+13,  1.96754595e+13,  1.25028059e+13,\n",
              "         4.67957067e+12, -1.66926107e+12, -3.93910042e+13,  4.37838181e+12,\n",
              "        -2.57189350e+12, -1.79522984e+13]),\n",
              " array([-2.41301935e+13, -1.20448597e+11, -1.68941001e+13,  2.65003396e+13,\n",
              "         1.28595360e+13,  8.44684136e+12,  5.81928428e+12, -8.33082171e+11,\n",
              "        -2.26809126e+13, -7.31588748e+12]),\n",
              " array([-9.21296694e+12,  5.90588642e+12, -1.81173975e+13, -1.93870552e+13,\n",
              "         3.06631161e+12, -8.45803664e+12, -5.81699005e+12, -3.65904663e+13,\n",
              "        -1.61098108e+13, -5.19614917e+12]),\n",
              " array([-2.73949405e+13, -2.40916031e+11, -6.72359323e+12,  3.08960108e+12,\n",
              "        -3.36267793e+13, -4.76926349e+12,  7.16268262e+12,  8.88709631e+12,\n",
              "        -3.22669853e+13, -3.55404469e+12]),\n",
              " array([ 7.87444500e+12, -2.91841289e+13, -1.05293702e+12,  1.18463598e+13,\n",
              "         8.97031918e+11, -2.43675668e+13, -8.05413653e+12,  8.36571244e+12,\n",
              "        -2.07135424e+13,  1.01648719e+13]),\n",
              " array([ 1.44048017e+13, -2.06911907e+13, -1.84155347e+12, -6.53507682e+11,\n",
              "        -5.56629123e+12, -2.27465492e+12, -2.26440724e+13,  1.36844315e+13,\n",
              "        -1.05570501e+12,  2.39242556e+13]),\n",
              " array([-3.83352328e+12, -1.00973628e+13, -1.66503831e+13, -5.78722781e+12,\n",
              "        -2.43879839e+13, -2.62670152e+13,  4.28550834e+12, -2.57129529e+13,\n",
              "        -1.41712325e+13,  3.99904555e+10]),\n",
              " array([-7.15555513e+12,  1.58612566e+12,  8.96705679e+12, -2.56708254e+13,\n",
              "        -3.38760428e+11,  1.28079596e+13, -1.45591444e+13,  1.25008292e+13,\n",
              "        -2.18865112e+13, -7.10745054e+12]),\n",
              " array([ 6.38532956e+12, -4.20125692e+13, -1.38145426e+13,  1.38009500e+13,\n",
              "         1.08941295e+13, -2.11479479e+13, -2.94407894e+12, -6.73743765e+11,\n",
              "        -5.58823539e+12, -2.04283741e+13]),\n",
              " array([-1.87934653e+13,  1.80105222e+13, -7.52061289e+12,  2.98324643e+12,\n",
              "        -9.52624655e+12, -1.93548917e+13, -6.08946593e+12, -2.55092860e+13,\n",
              "        -3.77184597e+12, -3.38779264e+13]),\n",
              " array([-1.70231594e+13, -3.89063982e+13, -2.87140121e+13, -2.45098851e+13,\n",
              "         8.30845465e+12, -5.41666912e+12, -1.45591702e+12, -1.06019755e+13,\n",
              "        -2.34304710e+13, -4.77426297e+12])]"
            ]
          },
          "metadata": {},
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "y_preds  # with layer_index=18"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "my_hlGB8bZV4",
        "outputId": "3eb85747-5527-43a8-de47-a8e84aacfd40"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[array([0.96408858, 0.88510712, 0.74801715, 0.41726325, 0.14450293,\n",
              "        0.94204061, 0.68642963, 0.23873456, 0.18076492, 0.89925655]),\n",
              " array([0.61956511, 0.82914544, 0.39426974, 0.74801086, 0.1350881 ,\n",
              "        0.6596291 , 0.53251797, 0.48623539, 0.9467238 , 0.32439335]),\n",
              " array([0.85825694, 0.62147529, 0.76711667, 0.95234358, 0.45985393,\n",
              "        0.63305692, 0.88612779, 0.49248477, 0.44888753, 0.21766493]),\n",
              " array([0.22425999, 0.30398587, 0.23830638, 0.382821  , 0.94705546,\n",
              "        0.56443306, 0.62120502, 0.54131019, 0.07111339, 0.45079669]),\n",
              " array([0.73627888, 0.11213173, 0.32320059, 0.09850259, 0.17090231,\n",
              "        0.05302753, 0.66926856, 0.92054819, 0.04092115, 0.0649834 ]),\n",
              " array([0.13063946, 0.95871142, 0.73914542, 0.29915338, 0.08589331,\n",
              "        0.72396739, 0.67769806, 0.10211055, 0.11268701, 0.76155581]),\n",
              " array([0.87900742, 0.32301346, 0.88566844, 0.46056626, 0.95554547,\n",
              "        0.18102891, 0.64117329, 0.43251788, 0.56286015, 0.09226081]),\n",
              " array([0.94012606, 0.67724452, 0.75702618, 0.08776972, 0.24547803,\n",
              "        0.62178128, 0.80347349, 0.94627612, 0.61228386, 0.33847219]),\n",
              " array([0.75494802, 0.31622798, 0.06194157, 0.68021156, 0.71375447,\n",
              "        0.82980007, 0.04783573, 0.47362228, 0.54524357, 0.4069766 ]),\n",
              " array([0.15799008, 0.1653532 , 0.32277415, 0.6312487 , 0.88012568,\n",
              "        0.1518761 , 0.89066168, 0.6904463 , 0.60949221, 0.38636204]),\n",
              " array([0.37255302, 0.92177151, 0.53549769, 0.47453843, 0.77912239,\n",
              "        0.56450529, 0.79433083, 0.60849753, 0.60538845, 0.28714437]),\n",
              " array([0.89255664, 0.52325501, 0.50082136, 0.91955523, 0.47731628,\n",
              "        0.36979005, 0.5064453 , 0.16431924, 0.02353021, 0.15864269]),\n",
              " array([0.52531402, 0.92018025, 0.20924646, 0.92902221, 0.1887657 ,\n",
              "        0.65579905, 0.3790168 , 0.84722156, 0.80683657, 0.51983535]),\n",
              " array([0.7054834 , 0.3242212 , 0.68165525, 0.86995792, 0.13977037,\n",
              "        0.04540505, 0.24366406, 0.48275323, 0.72900021, 0.20750097]),\n",
              " array([0.0253839 , 0.41708637, 0.4480164 , 0.08870838, 0.00907016,\n",
              "        0.72039994, 0.06477139, 0.93166007, 0.47912221, 0.64781735]),\n",
              " array([0.57860855, 0.58915784, 0.69957597, 0.67473547, 0.1719417 ,\n",
              "        0.84625658, 0.51753247, 0.71460117, 0.42590268, 0.46820958]),\n",
              " array([0.68955221, 0.07034647, 0.84991516, 0.47053583, 0.1750162 ,\n",
              "        0.56078   , 0.25433573, 0.27044004, 0.2731027 , 0.78911692]),\n",
              " array([0.47593949, 0.88916028, 0.32771839, 1.00495755, 0.66813444,\n",
              "        0.1955942 , 0.33173913, 0.81399721, 0.71440242, 0.68158695]),\n",
              " array([0.70096843, 0.66039507, 0.78801873, 0.6270835 , 0.2968498 ,\n",
              "        0.17250027, 0.26390376, 0.8075912 , 0.19128379, 0.7057339 ]),\n",
              " array([0.79256454, 0.93000319, 0.01191424, 0.06251566, 0.12764364,\n",
              "        0.19443186, 0.76752863, 0.31067844, 0.30677942, 0.04636196]),\n",
              " array([0.61693478, 0.88367643, 0.42764424, 0.5517247 , 0.71242501,\n",
              "        0.92888686, 0.52812318, 0.15887759, 0.43410016, 0.82203875]),\n",
              " array([0.06164778, 0.08818302, 0.16752998, 0.99603601, 0.59573963,\n",
              "        0.47077633, 0.14094381, 0.77269823, 0.78783704, 0.53975581]),\n",
              " array([1.00566188, 0.64581005, 0.23423067, 0.50183291, 0.13503312,\n",
              "        0.647415  , 0.46653627, 0.06941886, 0.29006766, 0.16162474]),\n",
              " array([0.33912802, 0.23804488, 0.20420201, 0.51269978, 0.98328175,\n",
              "        0.12871547, 0.05133293, 0.92177443, 0.78843663, 0.91969181]),\n",
              " array([0.37517199, 0.8480993 , 0.20905882, 0.69662054, 0.44590814,\n",
              "        0.05626704, 0.4526909 , 0.03871839, 0.82202589, 0.74402665]),\n",
              " array([0.0570261 , 0.24431423, 0.77490055, 0.16844588, 0.21732114,\n",
              "        0.28013222, 0.03252802, 0.24859909, 0.04470919, 0.9077881 ]),\n",
              " array([0.30793197, 0.10646063, 0.25245827, 0.10227304, 0.27430456,\n",
              "        0.67153183, 0.36117294, 0.19402333, 0.19236885, 0.48328431]),\n",
              " array([0.60996338, 0.54037736, 0.35094898, 0.51898918, 0.30822489,\n",
              "        0.81205241, 0.06178842, 0.57686748, 0.95117267, 0.78513639]),\n",
              " array([0.42240021, 0.67260976, 0.41859962, 0.14670325, 0.22084757,\n",
              "        0.7691956 , 0.44565871, 0.30514286, 0.16138277, 0.41131955]),\n",
              " array([0.43397371, 0.51957623, 0.91211414, 0.32220372, 0.89605186,\n",
              "        0.56708705, 0.84150718, 0.74884398, 0.91626542, 0.64239976]),\n",
              " array([0.38945997, 0.28298513, 0.08702682, 0.19610171, 0.06966381,\n",
              "        0.65211334, 0.46875482, 0.36718626, 0.67198555, 0.29369596]),\n",
              " array([ 0.58557873,  0.20166361,  0.39787273,  0.35192714,  0.68753521,\n",
              "         0.3398595 ,  0.49115452,  0.85416942,  0.62266473, -0.03033936]),\n",
              " array([0.19712577, 0.15775424, 0.96727365, 0.62828256, 0.5381514 ,\n",
              "        0.59163935, 0.40143777, 0.77489466, 0.81497732, 0.9702447 ]),\n",
              " array([ 0.29508379,  0.90252067,  0.57575493,  0.227061  ,  0.52647882,\n",
              "        -0.0094738 ,  0.87145314,  0.67060858,  0.39104077,  0.04929178]),\n",
              " array([0.84539891, 0.8721166 , 0.53687751, 0.57374375, 0.14954663,\n",
              "        0.35141526, 0.48667329, 0.33142719, 0.61328852, 0.481036  ]),\n",
              " array([ 0.85405969,  0.39290953,  0.60046964, -0.00255443,  0.27689486,\n",
              "         0.95764092,  0.62060498,  0.48438963,  0.14912792,  0.10071028]),\n",
              " array([0.3349037 , 0.61979423, 0.78296047, 0.3310923 , 0.84001595,\n",
              "        0.0616684 , 0.08645513, 0.03002371, 0.83677284, 0.36116214]),\n",
              " array([0.04907689, 0.10132951, 0.63829862, 0.01827266, 0.03727884,\n",
              "        0.07785855, 0.12199023, 0.42862718, 0.3949417 , 0.40530008]),\n",
              " array([ 0.81236794,  0.19307226,  0.45944711,  0.71807176,  0.02098221,\n",
              "         0.04445993,  0.3571132 , -0.00167954,  0.7226621 ,  0.61269043]),\n",
              " array([0.41567431, 0.44537187, 0.29261792, 0.57684408, 0.37218851,\n",
              "        0.26977872, 0.67188224, 0.37782138, 0.99082495, 0.13910716]),\n",
              " array([0.25534281, 0.08821304, 0.95650089, 0.22719752, 0.89448592,\n",
              "        0.07057899, 0.12650422, 0.85278876, 0.69694706, 0.06741974]),\n",
              " array([0.2616235 , 0.77139378, 0.39675481, 0.04545157, 0.4169178 ,\n",
              "        0.78188428, 0.3185542 , 0.95431963, 0.1361781 , 0.49338765]),\n",
              " array([0.87993223, 0.8802965 , 0.47684936, 0.90842782, 0.83582821,\n",
              "        0.30496346, 0.41945123, 0.80223014, 0.03234918, 0.56900408]),\n",
              " array([0.42911363, 0.33421178, 0.91405096, 0.46208082, 0.29914459,\n",
              "        0.90361807, 0.67715315, 0.05058124, 0.26540848, 0.21491847]),\n",
              " array([0.49008435, 0.80785688, 0.48122862, 0.34916025, 0.089261  ,\n",
              "        0.09567563, 0.07423366, 0.75684566, 0.29409824, 0.2862697 ]),\n",
              " array([0.54230976, 0.14957351, 0.66096957, 0.78405105, 0.55210188,\n",
              "        0.1610224 , 0.29389712, 0.57404575, 0.34810741, 0.76205953]),\n",
              " array([0.86690555, 0.83026117, 0.85405861, 0.68661587, 0.57334334,\n",
              "        0.3661496 , 0.07046105, 0.26554795, 0.61387401, 0.57268249]),\n",
              " array([0.86970746, 0.36519723, 0.30930465, 0.89742047, 0.56041868,\n",
              "        0.02449697, 0.38124279, 0.25406102, 0.15635416, 0.10472259]),\n",
              " array([0.86446088, 0.79129532, 0.22524829, 0.11197528, 0.62692045,\n",
              "        0.6284841 , 0.81134006, 0.45749698, 0.16668885, 0.68723006]),\n",
              " array([0.23034758, 0.6999672 , 0.86365037, 0.90946455, 0.76918508,\n",
              "        0.53756012, 0.70688806, 0.39209862, 0.64311802, 0.0814335 ]),\n",
              " array([0.96957002, 0.92524048, 0.3926017 , 0.2517493 , 0.3735987 ,\n",
              "        0.17675563, 0.06663724, 0.12817844, 0.92587898, 0.48606997]),\n",
              " array([0.45860289, 0.982563  , 0.27891345, 0.59343919, 0.35153987,\n",
              "        0.76935485, 0.20804236, 0.1887779 , 0.00994901, 0.53883659]),\n",
              " array([0.76338169, 0.17966523, 0.35060414, 0.18186314, 0.44383467,\n",
              "        0.54624256, 0.1068724 , 0.10665965, 0.33396428, 0.21393009]),\n",
              " array([0.48751215, 0.36006181, 0.13184282, 0.800293  , 0.80108485,\n",
              "        0.66638773, 0.21392332, 0.45700001, 0.68545959, 0.95919558]),\n",
              " array([0.78100607, 0.04525421, 0.80516217, 0.53147121, 0.95216979,\n",
              "        0.3579199 , 0.60601516, 0.92626134, 0.41971153, 0.93519085]),\n",
              " array([0.96032812, 0.18596206, 0.44465699, 0.58035108, 0.72952312,\n",
              "        0.32910966, 0.68708218, 0.77566239, 0.66268578, 0.86891551]),\n",
              " array([0.17141408, 0.51811202, 0.33881617, 0.22041343, 0.62876201,\n",
              "        0.2584693 , 0.89522682, 0.27228424, 0.58312852, 0.86134389]),\n",
              " array([0.21996462, 0.63027023, 0.41384629, 0.0121379 , 0.88493513,\n",
              "        0.14276418, 0.90553091, 0.72749981, 0.33299595, 0.68922746]),\n",
              " array([0.35739662, 0.25874996, 0.18087703, 0.17213501, 0.10870664,\n",
              "        0.29834817, 0.35901968, 0.79249472, 0.61045644, 0.95936265]),\n",
              " array([0.70428417, 0.90968615, 0.18764248, 0.87263606, 0.12890289,\n",
              "        0.08861998, 0.12491926, 0.95872414, 0.03461676, 0.11597786]),\n",
              " array([0.07844662, 0.76094586, 0.99434634, 0.67169068, 0.62669165,\n",
              "        0.80048172, 0.44686871, 1.00491656, 0.45790009, 0.32530961]),\n",
              " array([0.72230837, 0.59668313, 0.25298558, 0.91584585, 0.67024867,\n",
              "        0.13161412, 0.96585236, 0.49318833, 0.25128008, 0.66697121]),\n",
              " array([0.99420172, 0.77674392, 0.14892989, 0.81723158, 0.40650504,\n",
              "        0.12979875, 0.75411704, 0.44473257, 0.54730983, 0.90040197]),\n",
              " array([0.68096317, 0.42206547, 0.40484296, 0.89680821, 0.35367128,\n",
              "        0.51234704, 0.66892344, 0.45303283, 0.78319717, 0.41023951]),\n",
              " array([ 0.12316616,  0.21947583,  0.51537502,  0.35462196, -0.00351413,\n",
              "         0.34288466,  0.61757578,  0.33331037,  0.31226664,  0.4204751 ]),\n",
              " array([0.98788566, 0.20898976, 0.69902134, 0.52693485, 0.46914328,\n",
              "        0.91307197, 0.88422859, 0.38187318, 0.13284598, 0.66470212]),\n",
              " array([0.57145307, 0.82099733, 0.29046794, 0.17679622, 0.55918035,\n",
              "        0.85734486, 0.62562859, 0.95445689, 0.53369065, 0.15223062]),\n",
              " array([0.13956789, 0.7279719 , 0.71574456, 0.7712831 , 0.2341237 ,\n",
              "        0.07197511, 0.55565774, 0.74647736, 0.68406091, 0.56271101]),\n",
              " array([ 0.69033319,  0.85567437,  0.96716397,  0.90602508,  0.06080961,\n",
              "         0.00605643, -0.02491329,  0.21481928,  0.0929826 ,  0.15657535]),\n",
              " array([0.39544987, 0.32272877, 0.07472151, 0.10518614, 0.48565379,\n",
              "        0.19013136, 0.82444687, 0.05497717, 0.68875676, 0.82591615]),\n",
              " array([0.91581968, 0.68392846, 0.51938301, 1.01747325, 0.54228887,\n",
              "        0.14732918, 0.63443691, 0.78920046, 0.01215172, 0.78419448]),\n",
              " array([0.91931341, 0.01512914, 0.29097251, 0.34393254, 0.51829681,\n",
              "        0.94327141, 0.21499453, 0.84850604, 0.51010288, 0.83415815]),\n",
              " array([0.72831043, 0.6639631 , 0.61425673, 0.48933361, 0.82964447,\n",
              "        0.20182542, 0.8164179 , 0.77824941, 0.86918326, 0.31427861]),\n",
              " array([0.5716283 , 0.72308983, 0.24122883, 0.67016264, 0.90859055,\n",
              "        0.63150874, 0.63328872, 0.27758255, 0.10711568, 0.92300776]),\n",
              " array([0.49785917, 0.91995829, 0.31861626, 0.55601365, 0.35088742,\n",
              "        0.56763016, 0.55824012, 0.12404733, 0.59214736, 0.25149316]),\n",
              " array([0.94794863, 0.54647685, 0.86135618, 0.02195538, 0.6443497 ,\n",
              "        0.61474499, 0.07375101, 0.51603304, 0.84979562, 0.27926937]),\n",
              " array([0.0597703 , 0.15469452, 0.25103282, 0.61577675, 0.73927235,\n",
              "        0.96145549, 0.75166024, 0.60586994, 0.88586585, 0.49673368]),\n",
              " array([0.03528498, 0.92442308, 0.18523493, 0.38856202, 0.60849416,\n",
              "        0.76708957, 0.54764397, 0.29984349, 0.63078628, 0.81403454]),\n",
              " array([0.33907079, 0.67927319, 0.21692307, 0.30971636, 0.75747234,\n",
              "        0.61017648, 0.71473496, 0.56469611, 0.90934315, 0.23409219]),\n",
              " array([0.31025462, 0.95680055, 0.64482004, 0.53736053, 0.32771674,\n",
              "        0.64452713, 0.15902474, 0.60874135, 0.25666132, 0.60121664]),\n",
              " array([ 0.30632616,  0.05995998,  0.32785396,  0.86998955,  0.34815974,\n",
              "         0.70665476,  0.15444995, -0.00087923,  0.8741447 ,  0.47339457]),\n",
              " array([0.92270011, 0.94543088, 0.66747744, 0.09795946, 0.06850095,\n",
              "        0.93582255, 0.48214565, 0.22167103, 0.31774878, 0.8025264 ]),\n",
              " array([0.05856086, 0.4791231 , 0.12994381, 0.32132186, 0.16365109,\n",
              "        0.67531635, 0.69030402, 0.17002435, 0.52105634, 0.60614615]),\n",
              " array([0.43582498, 0.11680541, 0.58841722, 0.79685977, 0.8764205 ,\n",
              "        0.02174632, 0.83438697, 0.62717839, 0.08239456, 0.72948735]),\n",
              " array([0.4276752 , 0.22971789, 0.39873848, 0.57037943, 0.64058567,\n",
              "        0.26665312, 0.14936243, 0.38600285, 0.84141124, 0.47102758]),\n",
              " array([ 0.93496561,  0.59129828,  0.08051663, -0.0103892 ,  0.99581272,\n",
              "         0.46813445,  0.24510772,  0.48202612,  0.61715027,  0.33856267]),\n",
              " array([0.51447537, 0.40718393, 0.73404786, 0.204732  , 0.74888015,\n",
              "        0.69390643, 0.31735622, 0.2611568 , 0.03741201, 0.3012536 ]),\n",
              " array([0.94566843, 0.07023069, 0.83289913, 0.72569748, 0.84182704,\n",
              "        0.22138651, 0.92021955, 0.86943927, 0.42833968, 0.2279098 ]),\n",
              " array([0.50185368, 0.47764276, 0.78331541, 0.90829142, 0.29084022,\n",
              "        0.94857498, 0.43041843, 0.30287498, 0.36265718, 0.49665483]),\n",
              " array([0.86520779, 0.15486212, 0.41216775, 0.24804551, 0.81352763,\n",
              "        0.55885379, 0.25341145, 0.98354717, 0.79105345, 0.52916146])]"
            ]
          },
          "metadata": {},
          "execution_count": 22
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "RCUlzJ7heyV5"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}