{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-06-12T19:40:22.526813Z","iopub.execute_input":"2021-06-12T19:40:22.527164Z","iopub.status.idle":"2021-06-12T19:41:10.274129Z","shell.execute_reply.started":"2021-06-12T19:40:22.527088Z","shell.execute_reply":"2021-06-12T19:41:10.273332Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"#### Import Packages","metadata":{}},{"cell_type":"code","source":"%matplotlib inline\n\nimport os\nimport shutil\nimport random\nimport torch\nimport torchvision\nimport numpy as np\nimport time\nimport torch\nfrom torchvision import transforms, datasets\nfrom torch.utils.data import DataLoader, random_split\nimport pdb\nfrom PIL import Image\nfrom matplotlib import pyplot as plt\nimport argparse\ntorch.manual_seed(0)\n\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nprint('Using PyTorch version', torch.__version__)\n\n\nfrom os.path import isfile\nfrom os import rename\nSMOOTH=1\nimport pdb\nfrom sklearn.metrics import auc, roc_curve\nfrom PIL.ImageFilter import GaussianBlur\n\nimport pdb\nfrom numpy import pi as PI\nfrom numpy import sqrt\nfrom scipy.special import comb\n\n# tensorflow libraries\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom keras.preprocessing.image import ImageDataGenerator","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:10.276812Z","iopub.execute_input":"2021-06-12T19:41:10.277138Z","iopub.status.idle":"2021-06-12T19:41:17.380904Z","shell.execute_reply.started":"2021-06-12T19:41:10.277109Z","shell.execute_reply":"2021-06-12T19:41:17.380039Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nfrom sklearn.model_selection import train_test_split\nfrom pydicom import dcmread\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\n\n\nimport torch\nimport torch.nn as nn\nimport torchvision\nimport torchvision.transforms as transforms\nfrom torch.utils import data","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:17.382745Z","iopub.execute_input":"2021-06-12T19:41:17.383083Z","iopub.status.idle":"2021-06-12T19:41:17.602493Z","shell.execute_reply.started":"2021-06-12T19:41:17.383048Z","shell.execute_reply":"2021-06-12T19:41:17.601688Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"label_data = pd.read_csv('../input/rsna-pneumonia-detection-challenge/stage_2_train_labels.csv')\ncolumns = ['patientId', 'Target']\n\nlabel_data = label_data.filter(columns)\nlabel_data.head(5)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:17.605088Z","iopub.execute_input":"2021-06-12T19:41:17.605451Z","iopub.status.idle":"2021-06-12T19:41:17.691377Z","shell.execute_reply.started":"2021-06-12T19:41:17.605413Z","shell.execute_reply":"2021-06-12T19:41:17.689931Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_labels, val_labels = train_test_split(label_data.values, test_size=0.1)\nprint(train_labels.shape)\nprint(val_labels.shape)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:17.693222Z","iopub.execute_input":"2021-06-12T19:41:17.693926Z","iopub.status.idle":"2021-06-12T19:41:17.706182Z","shell.execute_reply.started":"2021-06-12T19:41:17.693885Z","shell.execute_reply":"2021-06-12T19:41:17.705084Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_f = '../input/rsna-pneumonia-detection-challenge/stage_2_train_images'\ntest_f = '../input/rsna-pneumonia-detection-challenge/stage_2_test_images'\n\ntrain_paths = [os.path.join(train_f, image[0]) for image in train_labels]\nval_paths = [os.path.join(train_f, image[0]) for image in val_labels]\n\nroot_dir = '../input/rsna-pneumonia-detection-challenge'\n\nprint(len(train_paths))\nprint(len(val_paths))","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:17.707821Z","iopub.execute_input":"2021-06-12T19:41:17.708189Z","iopub.status.idle":"2021-06-12T19:41:17.792307Z","shell.execute_reply.started":"2021-06-12T19:41:17.708152Z","shell.execute_reply":"2021-06-12T19:41:17.791226Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def imshow(num_to_show=9):\n    \n    plt.figure(figsize=(10,10))\n    \n    for i in range(num_to_show):\n        plt.subplot(3, 3, i+1)\n        plt.grid(False)\n        plt.xticks([])\n        plt.yticks([])\n        \n        img_dcm = dcmread(f'{train_paths[i+20]}.dcm')\n        img_np = img_dcm.pixel_array\n        plt.imshow(img_np, cmap=plt.cm.binary)\n        plt.xlabel(train_labels[i+20][1])\n\nimshow()","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:17.793698Z","iopub.execute_input":"2021-06-12T19:41:17.794102Z","iopub.status.idle":"2021-06-12T19:41:19.160629Z","shell.execute_reply.started":"2021-06-12T19:41:17.794060Z","shell.execute_reply":"2021-06-12T19:41:19.159878Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Dataset(data.Dataset):\n    \n    def __init__(self, paths, labels, transform=None):\n        self.paths = paths\n        self.labels = labels\n        self.transform = transform\n    \n    def __getitem__(self, index):\n        image = dcmread(f'{self.paths[index]}.dcm')\n        image = image.pixel_array\n        image = image / 255.0\n\n        image = (255*image).clip(0, 255).astype(np.uint8)\n        image = Image.fromarray(image).convert('L')\n\n        label = self.labels[index][1]\n        \n        if self.transform is not None:\n            image = self.transform(image)\n            \n        return image, label\n    \n    def __len__(self):\n        \n        return len(self.paths)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.163771Z","iopub.execute_input":"2021-06-12T19:41:19.164128Z","iopub.status.idle":"2021-06-12T19:41:19.171063Z","shell.execute_reply.started":"2021-06-12T19:41:19.164090Z","shell.execute_reply":"2021-06-12T19:41:19.169866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(train_paths)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.173495Z","iopub.execute_input":"2021-06-12T19:41:19.173866Z","iopub.status.idle":"2021-06-12T19:41:19.185725Z","shell.execute_reply.started":"2021-06-12T19:41:19.173811Z","shell.execute_reply":"2021-06-12T19:41:19.184715Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainset, testset = random_split(train_paths, [27000, 204])","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.187382Z","iopub.execute_input":"2021-06-12T19:41:19.187752Z","iopub.status.idle":"2021-06-12T19:41:19.212311Z","shell.execute_reply.started":"2021-06-12T19:41:19.187716Z","shell.execute_reply":"2021-06-12T19:41:19.211575Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Globally load device identifier\nos.environ['CUDA_VISIBLE_DEVICES'] = '0'\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.213500Z","iopub.execute_input":"2021-06-12T19:41:19.213847Z","iopub.status.idle":"2021-06-12T19:41:19.272663Z","shell.execute_reply.started":"2021-06-12T19:41:19.213799Z","shell.execute_reply":"2021-06-12T19:41:19.271700Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Miscellaneous initialization\ntorch.manual_seed(1)\nstart_time = time.time()\nparser = argparse.ArgumentParser()","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.274118Z","iopub.execute_input":"2021-06-12T19:41:19.274459Z","iopub.status.idle":"2021-06-12T19:41:19.285474Z","shell.execute_reply.started":"2021-06-12T19:41:19.274422Z","shell.execute_reply":"2021-06-12T19:41:19.284733Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"parser.add_argument('--num_epochs', type=int, default=5, help='Number of training epochs')\nparser.add_argument('--batch_size', type=int, default=32, help='Batch size')\nparser.add_argument('--lr', type=float, default=5e-4, help='Learning rate')\nparser.add_argument('--l2', type=float, default=0, help='L2 regularisation')\nparser.add_argument('--aug', action='store_true', default=False, help='Use data augmentation')\nparser.add_argument('--data_path', type=str, default=root_dir,help='Path to data.')\nparser.add_argument('--bond_dim', type=int, default=100, help='MPS Bond dimension')\nparser.add_argument('--nChannel', type=int, default=1, help='Number of input channels')\nparser.add_argument('--dense_net', action='store_true', default=False, help='Using Dense Net model')","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.286790Z","iopub.execute_input":"2021-06-12T19:41:19.287287Z","iopub.status.idle":"2021-06-12T19:41:19.299630Z","shell.execute_reply.started":"2021-06-12T19:41:19.287252Z","shell.execute_reply":"2021-06-12T19:41:19.298770Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"args = parser.parse_args([])","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.300906Z","iopub.execute_input":"2021-06-12T19:41:19.301486Z","iopub.status.idle":"2021-06-12T19:41:19.310699Z","shell.execute_reply.started":"2021-06-12T19:41:19.301446Z","shell.execute_reply":"2021-06-12T19:41:19.309918Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = args.batch_size\n\n# LoTeNet parameters\nadaptive_mode = False \nperiodic_bc   = False\n\nkernel = 2 # Stride along spatial dimensions\noutput_dim = 1 # output dimension\n \nfeature_dim = 2\n\n#logFile = time.strftime(\"%Y%m%d_%H_%M\")+'.txt'\n#makeLogFile(logFile)\n\nnormTensor = 0.5*torch.ones(args.nChannel)\n### Data processing and loading....","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.311908Z","iopub.execute_input":"2021-06-12T19:41:19.312286Z","iopub.status.idle":"2021-06-12T19:41:19.337647Z","shell.execute_reply.started":"2021-06-12T19:41:19.312247Z","shell.execute_reply":"2021-06-12T19:41:19.336961Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"### Data processing and loading....\ntrain_transform = transforms.Compose([transforms.RandomHorizontalFlip(),\n                                      transforms.Resize(size=(128,128)),\n                                      transforms.RandomVerticalFlip(),\n                                      transforms.RandomRotation(20),\n                                      transforms.ToTensor(),\n                                      transforms.Normalize(mean=normTensor,std=normTensor)])\n\nvalid_transform = torchvision.transforms.Compose([\n    torchvision.transforms.Resize(size=(128,128)),\n    torchvision.transforms.RandomHorizontalFlip(),\n    torchvision.transforms.RandomRotation(20),\n    torchvision.transforms.ToTensor(),\n    transforms.Normalize(mean=normTensor,std=normTensor)\n    #torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229,0.224,0.225])\n])","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.338772Z","iopub.execute_input":"2021-06-12T19:41:19.339145Z","iopub.status.idle":"2021-06-12T19:41:19.345713Z","shell.execute_reply.started":"2021-06-12T19:41:19.339103Z","shell.execute_reply":"2021-06-12T19:41:19.344699Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_dataset = Dataset(trainset, train_labels, transform=train_transform)\nval_dataset =  Dataset(val_paths, train_labels, transform=train_transform)\ntest_dataset = Dataset(testset, val_labels, transform=train_transform)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.347155Z","iopub.execute_input":"2021-06-12T19:41:19.347488Z","iopub.status.idle":"2021-06-12T19:41:19.356684Z","shell.execute_reply.started":"2021-06-12T19:41:19.347454Z","shell.execute_reply":"2021-06-12T19:41:19.355974Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_loader =  torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\ntest_loader =  torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\nval_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.359568Z","iopub.execute_input":"2021-06-12T19:41:19.359859Z","iopub.status.idle":"2021-06-12T19:41:19.368052Z","shell.execute_reply.started":"2021-06-12T19:41:19.359809Z","shell.execute_reply":"2021-06-12T19:41:19.367261Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def svd_flex(tensor, svd_string, max_D=None, cutoff=1e-10, sv_right=True,\n             sv_vec=None):\n    \"\"\"\n    Split an input tensor into two pieces using a SVD across some partition\n\n    Args:\n        tensor (Tensor):    Pytorch tensor with at least two indices\n\n        svd_string (str):   String of the form 'init_str->left_str,right_str',\n                            where init_str describes the indices of tensor, and\n                            left_str/right_str describe those of the left and\n                            right output tensors. The characters of left_str\n                            and right_str form a partition of the characters in\n                            init_str, but each contain one additional character\n                            representing the new bond which comes from the SVD\n\n                            Reversing the terms in svd_string to the left and\n                            right of '->' gives an ein_string which can be used\n                            to multiply both output tensors to give a (low rank\n                            approximation) of the input tensor\n\n        cutoff (float):     A truncation threshold which eliminates any\n                            singular values which are strictly less than cutoff\n\n        max_D (int):        A maximum allowed value for the new bond. If max_D\n                            is specified, the returned tensors\n\n        sv_right (bool):    The SVD gives two orthogonal matrices and a matrix\n                            of singular values. sv_right=True merges the SV\n                            matrix with the right output, while sv_right=False\n                            merges it with the left output\n\n        sv_vec (Tensor):    Pytorch vector with length max_D, which is modified\n                            in place to return the vector of singular values\n\n    Returns:\n        left_tensor (Tensor),\n        right_tensor (Tensor):  Tensors whose indices are described by the\n                                left_str and right_str parts of svd_string\n\n        bond_dim:               The dimension of the new bond appearing from\n                                the cutoff in our SVD. Note that this generally\n                                won't match the dimension of left_/right_tensor\n                                at this mode, which is padded with zeros\n                                whenever max_D is specified\n    \"\"\"\n    def prod(int_list):\n        output = 1\n        for num in int_list:\n            output *= num\n        return output\n\n    with torch.no_grad():\n        # Parse svd_string into init_str, left_str, and right_str\n        svd_string = svd_string.replace(' ', '')\n        init_str, post_str = svd_string.split('->')\n        left_str, right_str = post_str.split(',')\n\n        # Check formatting of init_str, left_str, and right_str\n        assert all([c.islower() for c in init_str+left_str+right_str])\n        assert len(set(init_str+left_str+right_str)) == len(init_str) + 1\n        assert len(set(init_str))+len(set(left_str))+len(set(right_str)) == \\\n               len(init_str)+len(left_str)+len(right_str)\n\n        # Get the special character representing our SVD-truncated bond\n        bond_char = set(left_str).intersection(set(right_str)).pop()\n        left_part = left_str.replace(bond_char, '')\n        right_part = right_str.replace(bond_char, '')\n\n        # Permute our tensor into something that can be viewed as a matrix\n        ein_str = f\"{init_str}->{left_part+right_part}\"\n        tensor = torch.einsum(ein_str, [tensor]).contiguous()\n\n        left_shape = list(tensor.shape[:len(left_part)])\n        right_shape = list(tensor.shape[len(left_part):])\n        left_dim, right_dim = prod(left_shape), prod(right_shape)\n\n        tensor = tensor.view([left_dim, right_dim])\n\n        # Get SVD and format so that left_mat * diag(svs) * right_mat = tensor\n        left_mat, svs, right_mat = torch.svd(tensor)\n        svs, _ = torch.sort(svs, descending=True)\n        right_mat = torch.t(right_mat)\n\n        # Decrease or increase our tensor sizes in the presence of max_D\n        if max_D and len(svs) > max_D:\n            svs = svs[:max_D]\n            left_mat = left_mat[:, :max_D]\n            right_mat = right_mat[:max_D]\n        elif max_D and len(svs) < max_D:\n            copy_svs = torch.zeros([max_D])\n            copy_svs[:len(svs)] = svs\n            copy_left = torch.zeros([left_mat.size(0), max_D])\n            copy_left[:, :left_mat.size(1)] = left_mat\n            copy_right = torch.zeros([max_D, right_mat.size(1)])\n            copy_right[:right_mat.size(0)] = right_mat\n            svs, left_mat, right_mat = copy_svs, copy_left, copy_right\n\n        # If given as input, copy singular values into sv_vec\n        if sv_vec is not None and svs.shape == sv_vec.shape:\n            sv_vec[:] = svs\n        elif sv_vec is not None and svs.shape != sv_vec.shape:\n            raise TypeError(f\"sv_vec.shape must be {list(svs.shape)}, but is \"\n                            f\"currently {list(sv_vec.shape)}\")\n\n        # Find the truncation point relative to our singular value cutoff\n        truncation = 0\n        for s in svs:\n            if s < cutoff:\n                break\n            truncation += 1\n        if truncation == 0:\n            raise RuntimeError(\"SVD cutoff too large, attempted to truncate \"\n                               \"tensor to bond dimension 0\")\n\n        # Perform the actual truncation\n        if max_D:\n            svs[truncation:] = 0\n            left_mat[:, truncation:] = 0\n            right_mat[truncation:] = 0\n        else:\n            # If max_D wasn't given, set it to the truncation index\n            max_D = truncation\n            svs = svs[:truncation]\n            left_mat = left_mat[:, :truncation]\n            right_mat = right_mat[:truncation]\n\n        # Merge the singular values into the appropriate matrix\n        if sv_right:\n            right_mat = torch.einsum('l,lr->lr', [svs, right_mat])\n        else:\n            left_mat = torch.einsum('lr,r->lr', [left_mat, svs])\n\n        # Reshape the matrices to make them proper tensors\n        left_tensor = left_mat.view(left_shape+[max_D])\n        right_tensor = right_mat.view([max_D]+right_shape)\n\n        # Finally, permute the indices into the desired order\n        if left_str != left_part + bond_char:\n            left_tensor = torch.einsum(f\"{left_part+bond_char}->{left_str}\",\n                                    [left_tensor])\n        if right_str != bond_char + right_part:\n            right_tensor = torch.einsum(f\"{bond_char+right_part}->{right_str}\",\n                                    [right_tensor])\n\n        return left_tensor, right_tensor, truncation\n\ndef init_tensor(shape, bond_str, init_method):\n    \"\"\"\n    Initialize a tensor with a given shape\n\n    Args:\n        shape:       The shape of our output parameter tensor.\n\n        bond_str:    The bond string describing our output parameter tensor,\n                     which is used in 'random_eye' initialization method.\n                     The characters 'l' and 'r' are used to refer to the\n                     left or right virtual indices of our tensor, and are\n                     both required to be present for the random_eye and\n                     min_random_eye initialization methods.\n\n        init_method: The method used to initialize the entries of our tensor.\n                     This can be either a string, or else a tuple whose first\n                     entry is an initialization method and whose remaining \n                     entries are specific to that method. In each case, std\n                     will always refer to a standard deviation for a random \n                     normal random component of each entry of the tensor. \n\n                     Allowed options are:\n                        * ('random_eye', std): Initialize each tensor input \n                            slice close to the identity\n                        * ('random_zero', std): Initialize each tensor input \n                            slice close to the zero matrix\n                        * ('min_random_eye', std, init_dim): Initialize each \n                            tensor input slice close to a truncated identity \n                            matrix, whose truncation leaves init_dim unit \n                            entries on the diagonal. If init_dim is larger\n                            than either of the bond dimensions, then init_dim \n                            is capped at the smaller bond dimension.\n    \"\"\"\n    # Unpack init_method if it is a tuple\n    if not isinstance(init_method, str):\n        init_str = init_method[0]\n        std = init_method[1]\n        if init_str == 'min_random_eye':\n            init_dim = init_method[2]\n\n        init_method = init_str\n    else:\n        std = 1e-9\n\n    # Check that bond_str is properly sized and doesn't have repeat indices\n    assert len(shape) == len(bond_str)\n    assert len(set(bond_str)) == len(bond_str)\n\n    if init_method not in ['random_eye', 'min_random_eye', 'random_zero']:\n        raise ValueError(f\"Unknown initialization method: {init_method}\")\n\n    if init_method in ['random_eye', 'min_random_eye']:\n        bond_chars = ['l', 'r']\n        assert all([c in bond_str for c in bond_chars])\n\n        # Initialize our tensor slices as identity matrices which each fill\n        # some or all of the initially allocated bond space\n        if init_method == 'min_random_eye':\n            \n            # The dimensions for our initial identity matrix. These will each \n            # be init_dim, unless init_dim exceeds one of the bond dimensions\n            bond_dims = [shape[bond_str.index(c)] for c in bond_chars]\n            if all([init_dim <= full_dim for full_dim in bond_dims]):\n                bond_dims = [init_dim, init_dim]\n            else:\n                init_dim = min(bond_dims)\n\n            eye_shape = [init_dim if c in bond_chars else 1 for c in bond_str]\n            expand_shape = [init_dim if c in bond_chars else shape[i]\n                            for i, c in enumerate(bond_str)]\n\n        elif init_method == 'random_eye':\n            eye_shape = [shape[i] if c in bond_chars else 1\n                         for i, c in enumerate(bond_str)]\n            expand_shape = shape\n            bond_dims = [shape[bond_str.index(c)] for c in bond_chars]\n\n        eye_tensor = torch.eye(bond_dims[0], bond_dims[1]).view(eye_shape)\n        eye_tensor = eye_tensor.expand(expand_shape)\n\n        tensor = torch.zeros(shape)\n        tensor[[slice(dim) for dim in expand_shape]] = eye_tensor\n\n        # Add on a bit of random noise\n        tensor += std * torch.randn(shape)\n\n    elif init_method == 'random_zero':\n        tensor = std * torch.randn(shape)\n\n    return tensor\n\n\n### OLDER MISCELLANEOUS FUNCTIONS ###\n\ndef onehot(labels, max_value):\n    \"\"\"\n    Convert a batch of labels from the set {0, 1,..., num_value-1} into their\n    onehot encoded counterparts\n    \"\"\"\n    label_vecs = torch.zeros([len(labels), max_value])\n\n    for i, label in enumerate(labels):\n        label_vecs[i, label] = 1.\n\n    return label_vecs\n\ndef joint_shuffle(input_data, input_labels):\n    \"\"\"\n    Shuffle input data and labels in a joint manner, so each label points to \n    its corresponding datum. Works for both regular and CUDA tensors\n    \"\"\"\n    assert input_data.is_cuda == input_labels.is_cuda\n    use_gpu = input_data.is_cuda\n    if use_gpu:\n        input_data, input_labels = input_data.cpu(), input_labels.cpu()\n\n    data, labels = input_data.numpy(), input_labels.numpy()\n\n    # Shuffle relative to the same seed\n    np.random.seed(0)\n    np.random.shuffle(data)\n    np.random.seed(0)\n    np.random.shuffle(labels)\n\n    data, labels = torch.from_numpy(data), torch.from_numpy(labels)\n    if use_gpu:\n        data, labels = data.cuda(), labels.cuda()\n\n    return data, labels\n\ndef load_HV_data(length):\n    \"\"\"\n    Output a toy \"horizontal/vertical\" data set of black and white\n    images with size length x length. Each image contains a single\n    horizontal or vertical stripe, set against a background\n    of the opposite color. The labels associated with these images\n    are either 0 (horizontal stripe) or 1 (vertical stripe).\n\n    In its current version, this returns two data sets, a training\n    set with 75% of the images and a test set with 25% of the\n    images.\n    \"\"\"\n    num_images = 4 * (2**(length-1) - 1)\n    num_patterns = num_images // 2\n    split = num_images // 4\n\n    if length > 14:\n        print(\"load_HV_data will generate {} images, \"\n              \"this could take a while...\".format(num_images))\n\n    images = np.empty([num_images,length,length], dtype=np.float32)\n    labels = np.empty(num_images, dtype=np.int)\n\n    # Used to generate the stripe pattern from integer i below\n    template = \"{:0\" + str(length) + \"b}\"\n\n    for i in range(1, num_patterns+1):\n        pattern = template.format(i)\n        pattern = [int(s) for s in pattern]\n\n        for j, val in enumerate(pattern):\n            # Horizontal stripe pattern\n            images[2*i-2, j, :] = val\n            # Vertical stripe pattern\n            images[2*i-1, :, j] = val\n\n        labels[2*i-2] = 0\n        labels[2*i-1] = 1\n\n    # Shuffle and partition into training and test sets\n    np.random.seed(0)\n    np.random.shuffle(images)\n    np.random.seed(0)\n    np.random.shuffle(labels)\n\n    train_images, train_labels = images[split:], labels[split:]\n    test_images, test_labels = images[:split], labels[:split]\n\n    return torch.from_numpy(train_images), \\\n           torch.from_numpy(train_labels), \\\n           torch.from_numpy(test_images), \\\n           torch.from_numpy(test_labels)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.369318Z","iopub.execute_input":"2021-06-12T19:41:19.369731Z","iopub.status.idle":"2021-06-12T19:41:19.415645Z","shell.execute_reply.started":"2021-06-12T19:41:19.369695Z","shell.execute_reply":"2021-06-12T19:41:19.413381Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def wCELoss(prediction, target):\n    w1 = 1.33  # False negative penalty\n    w2 = .66  # False positive penalty\n    return -torch.mean(w1 * target * torch.log(prediction.clamp_min(1e-3))\n    + w2 * (1. - target) * torch.log(1. - prediction.clamp_max(.999)))\n\nclass GaussianFilter(object):\n    \"\"\"Apply Gaussian blur to the PIL image\n    Args:\n    sigma (float): Sigma of Gaussian kernel. Default value 1.0\n    \"\"\"\n    def __init__(self, sigma=1):\n        self.sigma = sigma\n        self.filter = GaussianBlur(radius=sigma)\n        \n    def __call__(self, img):\n        \"\"\"\n        Args:\n            img (PIL Image): Image to be blurred.\n\n        Returns:\n            PIL Image: Blurred image.\n        \"\"\"\n        return img.filter(self.filter)\n\n    def __repr__(self):\n        return self.__class__.__name__ + '(sigma={})'.format(self.sigma)\n\nclass GaussianLayer(nn.Module):\n    def __init__(self):\n        super(GaussianLayer, self, sigma=1, size=10).__init__()\n        self.sigma = sigma\n        self.size = size\n        self.seq = nn.Sequential(\n            nn.ReflectionPad2d(size), \n            nn.Conv2d(3, 3, size, stride=1, padding=0, bias=None, groups=3)\n        )\n        self.weights_init()\n\n    def forward(self, x):\n        return self.seq(x)\n\n    def weights_init(self):\n        s = self.size * 2 + 1\n        k = np.zeros((s,s))\n        k[s,s] = 1\n        kernel = gaussian_filter(k,sigma=self.sigma)\n        for name, f in self.named_parameters():\n            f.data.copy_(torch.from_numpy(kernel))\n\nclass focalLoss(nn.Module):\n    def __init__(self, alpha=1, gamma=2, logits=False, reduce=True):\n        super(focalLoss, self).__init__()\n        self.alpha = alpha\n        self.gamma = gamma\n        self.logits = logits\n        self.reduce = reduce\n\n    def forward(self, inputs, targets):\n        if self.logits:\n            BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduce=False)\n        else:\n            BCE_loss = F.binary_cross_entropy(inputs, targets, reduce=False)\n        pt = torch.exp(-BCE_loss)\n        F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss\n\n        if self.reduce:\n            return torch.mean(F_loss)\n        else:\n            return F_loss\n\ndef computeAuc(target,preds):\n    fpr, tpr, thresholds = roc_curve(target,preds)\n    aucVal = auc(fpr,tpr)\n    return aucVal\n\nclass hingeLoss(torch.nn.Module):\n\n    def __init__(self):\n        super(hingeLoss, self).__init__()\n\n    def forward(self, output, target):\n#        pdb.set_trace()\n        target = 2*target-1\n        output = 2*output-1\n        hinge_loss = 1 - torch.mul(output, target)\n        hinge_loss[hinge_loss < 0] = 0\n        return hinge_loss.mean()\n\n\ndef makeBatchAdj(adj,bSize):\n\n    E = adj._nnz()\n    N = adj.shape[0]\n    batch_idx = torch.zeros(2,bSize*E).type(torch.LongTensor)\n    batch_val = torch.zeros(bSize*E)\n\n    idx = adj._indices()\n    vals = adj._values()\n\n    for i in range(bSize):\n        batch_idx[:,i*E:(i+1)*E] = idx + i*N\n        batch_val[i*E:(i+1)*E] = vals\n\n    return torch.sparse.FloatTensor(batch_idx,batch_val,(bSize*N,bSize*N))\n\n\ndef makeAdj(ngbrs, normalize=True):\n    \"\"\" Create an adjacency matrix, given the neighbour indices\n    Input: Nxd neighbourhood, where N is number of nodes\n    Output: NxN sparse torch adjacency matrix \n    \"\"\"\n#    pdb.set_trace()\n    N, d = ngbrs.shape\n    validNgbrs = (ngbrs >= 0) # Mask for valid neighbours amongst the d-neighbours\n    row = np.repeat(np.arange(N),d) # Row indices like in sparse matrix formats\n    row = row[validNgbrs.reshape(-1)] #Remove non-neighbour row indices \n    col = (ngbrs*validNgbrs).reshape(-1) # Obtain nieghbour col indices\n    col = col[validNgbrs.reshape(-1)] # Remove non-neighbour col indices\n    data = np.ones(col.size)\n    adj = sp.csr_matrix((np.ones(col.size, dtype=bool),(row, col)), shape=(N, N)).toarray() # Make adj matrix\n    adj = adj + np.eye(N) # Self connections \n    adj = sp.csr_matrix(adj, dtype=np.float32)#/(d+1)\n    if normalize:\n        adj = row_normalize(adj)\n    adj = sparse_mx_to_torch_sparse_tensor(adj) \n    \n    return adj\n\ndef makeRegAdj(numNgbrs=26):\n        \"\"\" Make regular pixel neighbourhoods\"\"\"\n        idx = 0\n        ngbrOffset = np.zeros((3,numNgbrs),dtype=int)\n        for i in range(-1,2):\n            for j in range(-1,2):\n                for k in range(-1,2):\n                        if(i | j | k):\n                                ngbrOffset[:,idx] = [i,j,k]\n                                idx+=1\n        idx = 0\n        ngbrs = np.zeros((numEl, numNgbrs), dtype=int)\n\n        for i in range(xdim):\n                for j in range(ydim):\n                        for k in range(zdim):\n                                xIdx = np.mod(ngbrOffset[0,:]+i,xdim)\n                                yIdx = np.mod(ngbrOffset[1,:]+j,ydim)\n                                zIdx = np.mod(ngbrOffset[2,:]+k,zdim)\n                                ngbrs[idx,:] = idxVol[xIdx, yIdx, zIdx]\n                                idx += 1\n\n\ndef makeAdjWithInvNgbrs(ngbrs, normalize=False):\n    \"\"\" Create an adjacency matrix, given the neighbour indices including invalid indices where self connections are added.\n    Input: Nxd neighbourhood, where N is number of nodes\n    Output: NxN sparse torch adjacency matrix \n    \"\"\"\n    np.random.seed(2)\n#    pdb.set_trace()\n    N, d = ngbrs.shape\n    row = np.arange(N).reshape(-1,1)\n    random = np.random.randint(0,N-1,(N,d))\n    valIdx = np.array((ngbrs < 0),dtype=int)\n    ngbrs = random*valIdx + ngbrs*(1-valIdx)# Mask for valid neighbours amongst the d-neighbours\n    row = np.repeat(row,d).reshape(-1) # Row indices like in sparse matrix formats\n    col = ngbrs.reshape(-1) # Obtain nieghbour col indices\n    data = np.ones(col.size)\n    adj = sp.csr_matrix((np.ones(col.size, dtype=bool),(row, col)), shape=(N, N)).toarray() # Make adj matrix\n    adj = adj + np.eye(N) # Self connections \n    adj = sp.csr_matrix(adj, dtype=np.float32)#/(d+1)\n    if normalize:\n        adj = row_normalize(adj)\n    adj = sparse_mx_to_torch_sparse_tensor(adj) \n    adj = adj.coalesce()\n    adj._values = adj.values()    \n    return adj\n\n\ndef transformers(adj):\n    \"\"\" Obtain source and sink node transformer matrices\"\"\"\n    edges = adj._indices()\n    N = adj.shape[0]\n    nnz = adj._nnz()\n    val = torch.ones(nnz)\n    idx0 = torch.arange(nnz)\n\n    idx = torch.stack((idx0,edges[1,:]))\n    n2e_in = torch.sparse.FloatTensor(idx,val,(nnz,N))\n\n    idx = torch.stack((idx0,edges[0,:]))\n    n2e_out = torch.sparse.FloatTensor(idx,val,(nnz,N))\n\n    return n2e_in, n2e_out\n\ndef sparse_mx_to_torch_sparse_tensor(sparse_mx):\n    \"\"\"Convert a scipy sparse matrix to a torch sparse tensor.\"\"\"\n    sparse_mx = sparse_mx.tocoo().astype(np.float32)\n    indices = torch.from_numpy(np.vstack((sparse_mx.row,\n                                          sparse_mx.col))).long()\n    values = torch.from_numpy(sparse_mx.data)\n    shape = torch.Size(sparse_mx.shape)\n    return torch.sparse.FloatTensor(indices, values, shape)\n\ndef to_linear_idx(x_idx, y_idx, num_cols):\n    assert num_cols > np.max(x_idx)\n    x_idx = np.array(x_idx, dtype=np.int32)\n    y_idx = np.array(y_idx, dtype=np.int32)\n    return y_idx * num_cols + x_idx\n\n\ndef row_normalize(mx):\n    \"\"\"Row-normalize sparse matrix\"\"\"\n    rowsum = np.array(mx.sum(1), dtype=np.float32)\n    r_inv = np.power(rowsum, -1).flatten()\n    r_inv[np.isinf(r_inv)] = 0.\n    r_mat_inv = sp.diags(r_inv)\n    mx = r_mat_inv.dot(mx)\n    return mx\n\ndef to_2d_idx(idx, num_cols):\n    idx = np.array(idx, dtype=np.int64)\n    y_idx = np.array(np.floor(idx / float(num_cols)), dtype=np.int64)\n    x_idx = idx % num_cols\n    return x_idx, y_idx\n\ndef dice_loss(preds, labels):\n    \"Return dice score. \"\n    preds_sq = preds**2\n    return 1 - (2. * (torch.sum(preds * labels)) + SMOOTH) / \\\n            (preds_sq.sum() + labels.sum() + SMOOTH)\n\ndef binary_accuracy(output, labels):\n    preds = output > 0.5\n    correct = preds.type_as(labels).eq(labels).double()\n    correct = correct.sum()\n    return correct / len(labels)\n\ndef multiClassAccuracy(output, labels):\n#    pdb.set_trace()\n    preds = output.argmax(1)\n#    preds = (output > (1.0/labels.shape[1])).type_as(labels)\n    correct = (preds == labels.view(-1))\n    correct = correct.sum().float()\n    return correct / len(labels)\n\ndef regrAcc(output, labels):\n#    pdb.set_trace()\n    preds = output.round().type(torch.long).type_as(labels)\n#    preds = (output > (1.0/labels.shape[1])).type_as(labels)\n    correct = (preds == labels.view(-1))\n    correct = correct.sum().float()\n    return correct / len(labels)\n\n\ndef rescaledRegAcc(output,labels,lRange=37,lMin=-20):\n#    pdb.set_trace()\n    preds = (output+1)*(lRange)/2 + lMin\n    preds = preds.round().type(torch.long).type_as(labels)\n#    preds = (output > (1.0/labels.shape[1])).type_as(labels)\n    correct = (preds == labels.view(-1))\n    correct = correct.sum().float()\n    return correct / len(labels)\n\ndef focalCE(preds, labels, gamma=1):\n    \"Return focal cross entropy\"\n    loss = -torch.mean( ( ((1-preds)**gamma) * labels * torch.log(preds) ) \\\n    + ( ((preds)**gamma) * (1-labels) * torch.log(1-preds) ) )\n    return loss\n\ndef dice(preds, labels):\n#    pdb.set_trace()\n    \"Return dice score\"\n    preds_bin = (preds > 0.5).type_as(labels)\n    return 2. * torch.sum(preds_bin * labels) / (preds_bin.sum() + labels.sum())\n\ndef wBCE(preds, labels, w):\n    \"Return weighted CE loss.\"\n    return -torch.mean( w*labels*torch.log(preds) + (1-w)*(1-labels)*torch.log(1-preds) )\n\ndef makeLogFile(filename=\"lossHistory.txt\"):\n    if isfile(filename):\n        rename(filename,\"lossHistoryOld.txt\")\n\n    with open(filename,\"w\") as text_file:\n        print('Epoch\\tlossTr\\taccTr\\tlossVl\\taccVl\\ttime(s)',file=text_file)\n    print(\"Log file created...\")\n    return\n\ndef writeLog(logFile, epoch, lossTr, accTr, lossVl, accVl,eTime):\n    print('Epoch:{:04d}\\t'.format(epoch + 1),\n          'lossTr:{:.4f}\\t'.format(lossTr),\n          'accTr:{:.4f}\\t'.format(accTr),\n          'lossVl:{:.4f}\\t'.format(lossVl),\n          'accVl:{:.4f}\\t'.format(accVl),\n          'time:{:.4f}'.format(eTime))\n\n    with open(logFile,\"a\") as text_file:\n        print('{:04d}\\t'.format(epoch + 1),\n                '{:.4f}\\t'.format(lossTr),\n                '{:.4f}\\t'.format(accTr),\n                '{:.4f}\\t'.format(lossVl),\n                '{:.4f}\\t'.format(accVl),\n                '{:.4f}'.format(eTime),file=text_file)\n    return\n\ndef plotLearningCurve():\n        plt.clf()\n        tmp = np.load('loss_tr.npz')['arr_0']\n        plt.plot(tmp,label='Tr.Loss')\n        tmp = np.load('loss_vl.npz')['arr_0']\n        plt.plot(tmp,label='Vl.Loss')\n        tmp = np.load('dice_tr.npz')['arr_0']\n        plt.plot(tmp,label='Tr.Dice')\n        tmp = np.load('dice_vl.npz')['arr_0']\n        plt.plot(tmp,label='Vl.Dice')\n        plt.legend()\n        plt.grid()\n        plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.417458Z","iopub.execute_input":"2021-06-12T19:41:19.418171Z","iopub.status.idle":"2021-06-12T19:41:19.501147Z","shell.execute_reply.started":"2021-06-12T19:41:19.418126Z","shell.execute_reply":"2021-06-12T19:41:19.498894Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Contractable:\n    \"\"\"\n    Container for tensors with labeled indices and a global batch size\n\n    The labels for our indices give some high-level knowledge of the tensor\n    layout, and permit the contraction of pairs of indices in a more \n    systematic manner. However, much of the actual heavy lifting is done \n    through specific contraction routines in different subclasses\n\n    Attributes:\n        tensor (Tensor):    A Pytorch tensor whose first index is a batch\n                            index. Sub-classes of Contractable may put other \n                            restrictions on tensor\n        bond_str (str):     A string whose letters each label a separate mode \n                            of our tensor, and whose length equals the order \n                            (number of modes) of our tensor\n        global_bs (int):    The batch size associated with all Contractables.\n                            This is shared between all Contractable instances \n                            and allows for automatic expanding of tensors\n    \"\"\"\n    # The global batch size\n    global_bs = None\n\n    def __init__(self, tensor, bond_str):\n        shape = list(tensor.shape)\n        num_dim = len(shape)\n        str_len = len(bond_str)\n\n        global_bs = Contractable.global_bs\n        batch_dim = tensor.size(0)\n\n        # Expand along a new batch dimension if needed\n        if ('b' not in bond_str and str_len == num_dim) or \\\n           ('b' == bond_str[0] and str_len == num_dim + 1):\n            if global_bs is not None:\n                tensor = tensor.unsqueeze(0).expand([global_bs] + shape)\n            else:\n                raise RuntimeError(\"No batch size given and no previous \"\n                                   \"batch size set\")\n            if bond_str[0] != 'b':\n                bond_str = 'b' + bond_str\n\n        # Check for correct formatting in bond_str\n        elif bond_str[0] != 'b' or str_len != num_dim:\n            raise ValueError(\"Length of bond string '{bond_str}' \"\n                            f\"({len(bond_str)}) must match order of \"\n                            f\"tensor ({len(shape)})\")\n\n        # Set the global batch size if it is unset or needs to be updated\n        elif global_bs is None or global_bs != batch_dim:\n            Contractable.global_bs = batch_dim\n\n        # Check that global batch size agrees with input tensor's first dim\n        elif global_bs != batch_dim:\n                raise RuntimeError(f\"Batch size previously set to {global_bs}\"\n                                    \", but input tensor has batch size \"\n                                   f\"{batch_dim}\")\n        \n        # Set the defining attributes of our Contractable\n        self.tensor = tensor\n        self.bond_str = bond_str\n\n    def __mul__(self, contractable, rmul=False):\n        \"\"\"\n        Multiply with another contractable along a linear index\n\n        The default behavior is to multiply the 'r' index of this instance\n        with the 'l' index of contractable, matching the batch ('b')\n        index of both, and take the outer product of other indices.\n        If rmul is True, contractable is instead multiplied on the right.\n        \"\"\"\n        # This method works for general Core subclasses besides Scalar (no 'l' \n        # and 'r' indices), composite contractables (no tensor attribute), and\n        # MatRegion (multiplication isn't just simple index contraction)\n        if isinstance(contractable, Scalar) or \\\n           not hasattr(contractable, 'tensor') or \\\n           type(contractable) is MatRegion:\n            return NotImplemented\n\n        tensors = [self.tensor, contractable.tensor]\n        bond_strs = [list(self.bond_str), list(contractable.bond_str)]\n        lowercases = [chr(c) for c in range(ord('a'), ord('z')+1)]\n\n        # Reverse the order of tensors if needed\n        if rmul:\n            tensors = tensors[::-1]\n            bond_strs = bond_strs[::-1]\n\n        # Check that bond strings are in proper format\n        for i, bs in enumerate(bond_strs):\n            assert bs[0] == 'b'\n            assert len(set(bs)) == len(bs)\n            assert all([c in lowercases for c in bs])\n            assert (i == 0 and 'r' in bs) or (i == 1 and 'l' in bs)\n\n        # Get used and free characters\n        used_chars = set(bond_strs[0]).union(bond_strs[1])\n        free_chars = [c for c in lowercases if c not in used_chars]\n\n        # Rename overlapping indices in the bond strings (except 'b', 'l', 'r')\n        specials = ['b', 'l', 'r']\n        for i, c in enumerate(bond_strs[1]):\n            if c in bond_strs[0] and c not in specials:\n                bond_strs[1][i] = free_chars.pop()\n\n        # Combine right bond of left tensor and left bond of right tensor\n        sum_char = free_chars.pop()\n        bond_strs[0][bond_strs[0].index('r')] = sum_char\n        bond_strs[1][bond_strs[1].index('l')] = sum_char\n        specials.append(sum_char)\n\n        # Build bond string of ouput tensor\n        out_str = ['b']\n        for bs in bond_strs:\n            out_str.extend([c for c in bs if c not in specials])\n        out_str.append('l' if 'l' in bond_strs[0] else '')\n        out_str.append('r' if 'r' in bond_strs[1] else '')\n\n        # Build the einsum string for this operation\n        bond_strs = [''.join(bs) for bs in bond_strs]\n        out_str = ''.join(out_str)\n        ein_str = f\"{bond_strs[0]},{bond_strs[1]}->{out_str}\"\n\n        # Contract along the linear dimension to get an output tensor\n        out_tensor = torch.einsum(ein_str, [tensors[0], tensors[1]])\n\n        # Return our output tensor wrapped in an appropriate class\n        if out_str == 'br':\n            return EdgeVec(out_tensor, is_left_vec=True)\n        elif out_str == 'bl':\n            return EdgeVec(out_tensor, is_left_vec=False)\n        elif out_str == 'blr':\n            return SingleMat(out_tensor)\n        elif out_str == 'bolr':\n            return OutputCore(out_tensor)\n        else:\n            return Contractable(out_tensor, out_str)\n\n    def __rmul__(self, contractable):\n        \"\"\"\n        Multiply with another contractable along a linear index\n        \"\"\"\n        return self.__mul__(contractable, rmul=True)\n\n    def reduce(self):\n        \"\"\"\n        Return the contractable without any modification\n\n        reduce() can be any method which returns a contractable. This is\n        trivially possible for any contractable by returning itself\n        \"\"\"\n        return self\n\nclass ContractableList(Contractable):\n    \"\"\"\n    A list of contractables which can all be multiplied together in order\n\n    Calling reduce on a ContractableList instance will first reduce every item\n    to a linear contractable, and then contract everything together\n    \"\"\"\n    def __init__(self, contractable_list):\n        # Check that input list is nonempty and has contractables as entries\n        if not isinstance(contractable_list, list) or contractable_list is []:\n            raise ValueError(\"Input to ContractableList must be nonempty list\")\n        for i, item in enumerate(contractable_list):\n            if not isinstance(item, Contractable):\n                raise ValueError(\"Input items to ContractableList must be \"\n                                f\"Contractable instances, but item {i} is not\")\n\n        self.contractable_list = contractable_list\n\n    def __mul__(self, contractable, rmul=False):\n        \"\"\"\n        Multiply a contractable by everything in ContractableList in order\n        \"\"\"\n        # The input cannot be a composite contractable\n        assert hasattr(contractable, 'tensor')\n        output = contractable.tensor\n\n        # Multiply by everything in ContractableList, in the correct order\n        if rmul:\n            for item in self.contractable_list:\n                output = item * output\n        else:\n            for item in self.contractable_list[::-1]:\n                output = output * item\n\n        return output\n\n    def __rmul__(self, contractable):\n        \"\"\"\n        Multiply another contractable by everything in ContractableList\n        \"\"\"\n        return self.__mul__(contractable, rmul=True)\n\n    def reduce(self, parallel_eval=False):\n        \"\"\"\n        Reduce all the contractables in list before multiplying them together\n        \"\"\"\n        c_list = self.contractable_list\n        # For parallel_eval, reduce all contractables in c_list\n        if parallel_eval:\n            c_list = [item.reduce() for item in c_list]\n\n        # Multiply together all the contractables. This multiplies in right to \n        # left order, but certain inefficient contractions are unsupported.\n        # If we encounter an unsupported operation, then try multiplying from\n        # the left end of the list instead\n        while len(c_list) > 1:\n            try:\n                c_list[-2] = c_list[-2] * c_list[-1]\n                del c_list[-1]\n            except TypeError:\n                c_list[1] = c_list[0] * c_list[1]\n                del c_list[0]\n\n        return c_list[0]\n\nclass MatRegion(Contractable):\n    \"\"\"\n    A contiguous collection of matrices which are multiplied together\n\n    The input tensor defining our MatRegion must have shape \n    [batch_size, num_mats, D, D], or [num_mats, D, D] when the global batch\n    size is already known\n    \"\"\"\n    def __init__(self, mats):\n        shape = list(mats.shape)\n        if len(shape) not in [3, 4] or shape[-2] != shape[-1]:\n            raise ValueError(\"MatRegion tensors must have shape \"\n                             \"[batch_size, num_mats, D, D], or [num_mats,\"\n                             \" D, D] if batch size has already been set\")\n\n        super().__init__(mats, bond_str='bslr')\n\n    def __mul__(self, edge_vec, rmul=False):\n        \"\"\"\n        Iteratively multiply an input vector with all matrices in MatRegion\n        \"\"\"\n        # The input must be an instance of EdgeVec\n        if not isinstance(edge_vec, EdgeVec):\n            return NotImplemented\n\n        mats = self.tensor\n        num_mats = mats.size(1)\n        batch_size = mats.size(0)\n\n        # Load our vector and matrix batches\n        dummy_ind = 1 if rmul else 2\n        vec = edge_vec.tensor.unsqueeze(dummy_ind)\n        mat_list = [mat.squeeze(1) for mat in torch.chunk(mats, num_mats, 1)]\n\n        # Do the repeated matrix-vector multiplications in the proper order\n        log_norm = 0\n        for i, mat in enumerate(mat_list[::(1 if rmul else -1)], 1):\n            if rmul:\n                vec = torch.bmm(vec, mat)\n            else:\n                vec = torch.bmm(mat, vec)\n\n        # Since we only have a single vector, wrap it as a EdgeVec\n        return EdgeVec(vec.squeeze(dummy_ind), is_left_vec=rmul)\n\n    def __rmul__(self, edge_vec):\n        return self.__mul__(edge_vec, rmul=True)\n\n    def reduce(self):\n        \"\"\"\n        Multiplies together all matrices and returns resultant SingleMat\n\n        This method uses iterated batch multiplication to evaluate the full \n        matrix product in depth O( log(num_mats) )\n        \"\"\"\n        mats = self.tensor\n        shape = list(mats.shape)\n        batch_size = mats.size(0)\n        size, D = shape[1:3]\n\n        # Iteratively multiply pairs of matrices until there is only one\n        while size > 1:\n            odd_size = (size % 2 == 1)\n            half_size = size // 2\n            nice_size = 2 * half_size\n        \n            even_mats = mats[:, 0:nice_size:2]\n            odd_mats = mats[:, 1:nice_size:2]\n            # For odd sizes, set aside one batch of matrices for the next round\n            leftover = mats[:, nice_size:]\n\n            # Multiply together all pairs of matrices (except leftovers)\n            mats = torch.einsum('bslu,bsur->bslr', [even_mats, odd_mats])\n            mats = torch.cat([mats, leftover], 1)\n\n            size = half_size + int(odd_size)\n\n        # Since we only have a single matrix, wrap it as a SingleMat\n        return SingleMat(mats.squeeze(1))\n\nclass OutputCore(Contractable):\n    \"\"\"\n    A single MPS core with a single output index\n    \"\"\"\n    def __init__(self, tensor):\n        # Check the input shape\n        if len(tensor.shape) not in [3, 4]:\n            raise ValueError(\"OutputCore tensors must have shape [batch_size, \"\n                             \"output_dim, D_l, D_r], or else [output_dim, D_l,\"\n                             \" D_r] if batch size has already been set\")\n\n        super().__init__(tensor, bond_str='bolr')\n\nclass SingleMat(Contractable):\n    \"\"\"\n    A batch of matrices associated with a single location in our MPS\n    \"\"\"\n    def __init__(self, mat):\n        # Check the input shape\n        if len(mat.shape) not in [2, 3]:\n            raise ValueError(\"SingleMat tensors must have shape [batch_size, \"\n                             \"D_l, D_r], or else [D_l, D_r] if batch size \"\n                             \"has already been set\")\n\n        super().__init__(mat, bond_str='blr')\n\nclass OutputMat(Contractable):\n    \"\"\"\n    An output core associated with an edge of our MPS\n    \"\"\"\n    def __init__(self, mat, is_left_mat):\n        # Check the input shape\n        if len(mat.shape) not in [2, 3]:\n            raise ValueError(\"OutputMat tensors must have shape [batch_size, \"\n                             \"D, output_dim], or else [D, output_dim] if \"\n                             \"batch size has already been set\")\n\n        # OutputMats on left edge will have a right-facing bond, and vice versa\n        bond_str = 'b' + ('r' if is_left_mat else 'l') + 'o'\n        super().__init__(mat, bond_str=bond_str)\n\n    def __mul__(self, edge_vec, rmul=False):\n        \"\"\"\n        Multiply with an edge vector along the shared linear index\n        \"\"\"\n        if not isinstance(edge_vec, EdgeVec):\n            raise NotImplemented\n        else:\n            return super().__mul__(edge_vec, rmul)\n\n    def __rmul__(self, edge_vec):\n        return self.__mul__(edge_vec, rmul=True)\n\nclass EdgeVec(Contractable):\n    \"\"\"\n    A batch of vectors associated with an edge of our MPS\n\n    EdgeVec instances are always associated with an edge of an MPS, which \n    requires the is_left_vec flag to be set to True (vector on left edge) or \n    False (vector on right edge)\n    \"\"\"\n    def __init__(self, vec, is_left_vec):\n        # Check the input shape\n        if len(vec.shape) not in [1, 2]:\n            raise ValueError(\"EdgeVec tensors must have shape \"\n                             \"[batch_size, D], or else [D] if batch size \"\n                             \"has already been set\")\n\n        # EdgeVecs on left edge will have a right-facing bond, and vice versa\n        bond_str = 'b' + ('r' if is_left_vec else 'l')\n        super().__init__(vec, bond_str=bond_str)\n\n    def __mul__(self, right_vec):\n        \"\"\"\n        Take the inner product of our vector with another vector\n        \"\"\"\n        # The input must be an instance of EdgeVec\n        if not isinstance(right_vec, EdgeVec):\n            return NotImplemented\n\n        left_vec = self.tensor.unsqueeze(1)\n        right_vec = right_vec.tensor.unsqueeze(2)\n        batch_size = left_vec.size(0)\n\n        # Do the batch inner product\n        scalar = torch.bmm(left_vec, right_vec).view([batch_size])\n\n        # Since we only have a single scalar, wrap it as a Scalar\n        return Scalar(scalar)\n\nclass Scalar(Contractable):\n    \"\"\"\n    A batch of scalars\n    \"\"\"\n    def __init__(self, scalar):\n        # Add dummy dimension if we have a torch scalar\n        shape = list(scalar.shape)\n        if shape is []:\n            scalar = scalar.view([1])\n            shape = [1]\n            \n        # Check the input shape\n        if len(shape) != 1:\n            raise ValueError(\"input scalar must be a torch tensor with shape \"\n                             \"[batch_size], or [] or [1] if batch size has \"\n                             \"been set\")\n\n        super().__init__(scalar, bond_str='b')\n\n    def __mul__(self, contractable):\n        \"\"\"\n        Multiply a contractable by our scalar and return the result\n        \"\"\"\n        scalar = self.tensor\n        tensor = contractable.tensor\n        bond_str = contractable.bond_str\n\n        ein_string = f\"{bond_str},b->{bond_str}\"\n        out_tensor = torch.einsum(ein_string, [tensor, scalar])\n\n        # Wrap the result in the same class right_contractable belongs to\n        contract_class = type(contractable)\n        if contract_class is not Contractable:\n            return contract_class(out_tensor)\n        else:\n            return Contractable(out_tensor, bond_str)\n\n    def __rmul__(self, contractable):\n        # Scalar multiplication is commutative\n        return self.__mul__(contractable)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.504229Z","iopub.execute_input":"2021-06-12T19:41:19.505365Z","iopub.status.idle":"2021-06-12T19:41:19.577959Z","shell.execute_reply.started":"2021-06-12T19:41:19.505322Z","shell.execute_reply":"2021-06-12T19:41:19.577230Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class MPS(nn.Module):\n\t\"\"\"\n\tMatrix product state which converts input into a single output vector\n\t\"\"\"\n\tdef __init__(self, input_dim, output_dim, bond_dim, feature_dim=2, nCh=3,\n\t\t\t\t adaptive_mode=False, periodic_bc=False, parallel_eval=False,\n\t\t\t\t label_site=None, path=None, init_std=1e-9, use_bias=True,\n\t\t\t\t fixed_bias=True, cutoff=1e-10, merge_threshold=2000):\n\t\tsuper().__init__()\n\n\t\torg_input_dim = input_dim\n\t\tif label_site is None:\n\t\t\tlabel_site = input_dim // 2\n\t\tassert label_site >= 0 and label_site <= input_dim\n\n\t\t# Using bias matrices in adaptive_mode is too complicated, so I'm \n\t\t# disabling it here\n\t\tif adaptive_mode:\n\t\t\tuse_bias = False\n\t\t# Our MPS is made of two InputRegions separated by an OutputSite.\n\t\tmodule_list = []\n\t\tinit_args = {'bond_str': 'slri',\n\t\t\t\t\t 'shape': [label_site, bond_dim, bond_dim, feature_dim],\n\t\t\t\t\t 'init_method': ('min_random_eye' if adaptive_mode else\n\t\t\t\t\t 'random_zero', init_std, output_dim)}\n\n\t\t# The first input region\n\t\tif label_site > 0:\n\t\t\ttensor = init_tensor(**init_args)\n\n\t\t\tmodule_list.append(InputRegion(tensor, use_bias=use_bias, \n\t\t\t\t\t\t\t\t\t\t   fixed_bias=fixed_bias))\n\n\t\t# The output site\n\t\ttensor = init_tensor(shape=[output_dim, bond_dim, bond_dim],\n\t\t\tbond_str='olr', init_method=('min_random_eye' if adaptive_mode else\n\t\t\t\t\t\t\t\t\t\t 'random_eye', init_std, output_dim))\n\t\tmodule_list.append(OutputSite(tensor))\n\n\t\t# The other input region\n\t\tif label_site < input_dim:\n\t\t\tinit_args['shape'] = [input_dim-label_site, bond_dim, bond_dim, \n\t\t\t\t\t\t\t\t  feature_dim]\n\t\t\ttensor = init_tensor(**init_args)\n\t\t\tmodule_list.append(InputRegion(tensor, use_bias=use_bias, \n\t\t\t\t\t\t\t\t\t\t   fixed_bias=fixed_bias))\n\n\t\t# Initialize linear_region according to our adaptive_mode specification\n\t\tif adaptive_mode:\n\t\t\tself.linear_region = MergedLinearRegion(module_list=module_list,\n\t\t\t\t\t\t\t\t periodic_bc=periodic_bc,\n\t\t\t\t\t\t\t\t parallel_eval=parallel_eval, cutoff=cutoff,\n\t\t\t\t\t\t\t\t merge_threshold=merge_threshold)\n\n\t\t\t# Initialize the list of bond dimensions, which starts out constant\n\t\t\tself.bond_list = bond_dim * torch.ones(input_dim + 2, \n\t\t\t\t\t\t\t\t\t\t\t\t   dtype=torch.long)\n\t\t\tif not periodic_bc:\n\t\t\t\tself.bond_list[0], self.bond_list[-1] = 1, 1\n\n\t\t\t# Initialize the list of singular values, which start out at -1\n\t\t\tself.sv_list = -1. * torch.ones([input_dim + 2, bond_dim])\n\n\t\telse:\n\t\t\tself.linear_region = LinearRegion(module_list=module_list,\n\t\t\t\t\t\t\t\t periodic_bc=periodic_bc,\n\t\t\t\t\t\t\t\t parallel_eval=parallel_eval)\n\t\tassert len(self.linear_region) == input_dim\n\n\t\tif path:\n\t\t\tassert isinstance(path, (list, torch.Tensor))\n\t\t\tassert len(path) == input_dim\n\n\t\t# Set the rest of our MPS attributes\n\t\tself.input_dim = input_dim\n\t\tself.output_dim = output_dim\n\t\tself.bond_dim = bond_dim\n\t\tself.feature_dim = feature_dim\n\t\tself.periodic_bc = periodic_bc\n\t\tself.adaptive_mode = adaptive_mode\n\t\tself.label_site = label_site\n\t\tself.path = path\n\t\tself.use_bias = use_bias\n\t\tself.fixed_bias = fixed_bias\n\t\tself.cutoff = cutoff\n\t\tself.merge_threshold = merge_threshold\n\t\tself.feature_map = None\n\t\tself.linear_region = self.linear_region.to(device)\n\t\tmodule_list = [m.to(device) for m in module_list]\n\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tEmbed our data and pass it to an MPS with a single output site\n\n\t\tArgs:\n\t\t\tinput_data (Tensor): Input with shape [batch_size, input_dim] or\n\t\t\t\t\t\t\t\t [batch_size, input_dim, feature_dim]. In the\n\t\t\t\t\t\t\t\t former case, the data points are turned into\n\t\t\t\t\t\t\t\t 2D vectors using a default linear feature map.\n\t\t\"\"\"\n\t\tinput_data = input_data.permute(0,2,1)\n\n\t\tx1 = torch.cos(input_data * PI/2)\n\t\tx2 = torch.sin(input_data* PI/2)\t\t\n\t\tx = torch.cat((x1,x2),dim=2)\n\n\t\toutput = self.linear_region(x) \n\n\t\treturn output.squeeze()\n\n\tdef core_len(self):\n\t\t\"\"\"\n\t\tReturns the number of cores, which is at least the required input size\n\t\t\"\"\"\n\t\treturn self.linear_region.core_len()\n\n\tdef __len__(self):\n\t\t\"\"\"\n\t\tReturns the number of input sites, which equals the input size\n\t\t\"\"\"\n\t\treturn self.input_dim\n\nclass LinearRegion(nn.Module):\n\t\"\"\"\n\tList of modules which feeds input to each module and returns reduced output\n\t\"\"\"\n\tdef __init__(self, module_list,periodic_bc=False, parallel_eval=False,\n\t\t\t\t module_states=None):\n\t\t# Check that module_list is a list whose entries are Pytorch modules\n\t\tif not isinstance(module_list, list) or module_list is []:\n\t\t\traise ValueError(\"Input to LinearRegion must be nonempty list\")\n\t\tfor i, item in enumerate(module_list):\n\t\t\tif not isinstance(item, nn.Module):\n\t\t\t\traise ValueError(\"Input items to LinearRegion must be PyTorch \"\n\t\t\t\t\t\t\t\tf\"Module instances, but item {i} is not\")\n\t\tsuper().__init__()\n\n\t\t# Wrap as a ModuleList for proper parameter registration\n\t\tself.module_list = nn.ModuleList(module_list)\n\t\tself.periodic_bc = periodic_bc\n\t\tself.parallel_eval = parallel_eval\n\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tContract input with list of MPS cores and return result as contractable\n\n\t\tArgs:\n\t\t\tinput_data (Tensor): Input with shape [batch_size, input_dim,\n\t\t\t\t\t\t\t\t\t\t\t\t   feature_dim]\n\t\t\"\"\"\n\t\t# Check that input_data has the correct shape\n\t\tassert len(input_data.shape) == 3\n\t\tassert input_data.size(1) == len(self)\n\t\tperiodic_bc = self.periodic_bc\n\t\tparallel_eval = self.parallel_eval\n\t\tlin_bonds = ['l', 'r']\n\n\t\t# For each module, pull out the number of pixels needed and call that\n\t\t# module's forward() method, putting the result in contractable_list\n\t\tind = 0\n\t\tcontractable_list = []\n\t\tfor module in self.module_list:\n\t\t\tmod_len = len(module)\n\t\t\tif mod_len == 1:\n\t\t\t\tmod_input = input_data[:, ind]\n\t\t\telse:\n\t\t\t\tmod_input = input_data[:, ind:(ind+mod_len)]\n\t\t\tind += mod_len\n\n\t\t\tcontractable_list.append(module(mod_input))\n\n\t\t# For periodic boundary conditions, reduce contractable_list and\n\t\t# trace over the left and right indices to get our output\n\t\tif periodic_bc:\n\t\t\tcontractable_list = ContractableList(contractable_list)\n\t\t\tcontractable = contractable_list.reduce(parallel_eval=True)\n\n\t\t\t# Unpack the output (atomic) contractable\n\t\t\ttensor, bond_str = contractable.tensor, contractable.bond_str\n\t\t\tassert all(c in bond_str for c in lin_bonds)\n\n\t\t\t# Build einsum string for the trace of tensor\n\t\t\tin_str, out_str = \"\", \"\"\n\t\t\tfor c in bond_str:\n\t\t\t\tif c in lin_bonds:\n\t\t\t\t\tin_str += 'l'\n\t\t\t\telse:\n\t\t\t\t\tin_str += c\n\t\t\t\t\tout_str += c\n\t\t\tein_str = in_str + \"->\" + out_str\n\n\t\t\t# Return the trace over left and right indices\n\t\t\treturn torch.einsum(ein_str, [tensor])\n\n\t\t# For open boundary conditions, add dummy edge vectors to\n\t\t# contractable_list and reduce everything to get our output\n\t\telse:\n\t\t\t# Get the dimension of left and right bond indices\n\t\t\tend_items = [contractable_list[i]for i in [0, -1]]\n\t\t\tbond_strs = [item.bond_str for item in end_items]\n\t\t\tbond_inds = [bs.index(c) for (bs, c) in zip(bond_strs, lin_bonds)]\n\t\t\tbond_dims = [item.tensor.size(ind) for (item, ind) in\n\t\t\t\t\t\t\t\t\t\t\t   zip(end_items, bond_inds)]\n\n\t\t\t# Build dummy end vectors and insert them at the ends of our list\n\t\t\tend_vecs = [torch.zeros(dim) for dim in bond_dims]\n\t\t\tend_vecs = [e.to(device) for e in end_vecs]\n\n\t\t\tfor vec in end_vecs:\n\t\t\t\tvec[0] = 1\n\n\t\t\tcontractable_list.insert(0, EdgeVec(end_vecs[0], is_left_vec=True))\n\t\t\tcontractable_list.append(EdgeVec(end_vecs[1], is_left_vec=False))\n\n\t\t\t# Multiply together everything in contractable_list\n\t\t\tcontractable_list = ContractableList(contractable_list)\n\t\t\toutput = contractable_list.reduce(parallel_eval=parallel_eval)\n\n\t\t\treturn output.tensor\n\n\tdef core_len(self):\n\t\t\"\"\"\n\t\tReturns the number of cores, which is at least the required input size\n\t\t\"\"\"\n\t\treturn sum([module.core_len() for module in self.module_list])\n\n\tdef __len__(self):\n\t\t\"\"\"\n\t\tReturns the number of input sites, which is the required input size\n\t\t\"\"\"\n\t\treturn sum([len(module) for module in self.module_list])\n\nclass MergedLinearRegion(LinearRegion):\n\t\"\"\"\n\tDynamic variant of LinearRegion that periodically rearranges its submodules\n\t\"\"\"\n\tdef __init__(self, module_list, periodic_bc=False, parallel_eval=False, \t\t\t\t cutoff=1e-10, merge_threshold=2000):\n\t\t# Initialize a LinearRegion with our given module_list\n\t\tsuper().__init__(module_list, periodic_bc, parallel_eval)\n\n\t\t# Initialize attributes self.module_list_0 and self.module_list_1\n\t\t# using the unmerged self.module_list, then redefine the latter in\n\t\t# terms of one of the former lists\n\t\tself.offset = 0\n\t\tself.merge(offset=self.offset)\n\t\tself.merge(offset=(self.offset+1)%2)\n\t\tself.module_list = getattr(self, f\"module_list_{self.offset}\")\n\n\t\t# Initialize variables used during switching\n\t\tself.input_counter = 0\n\t\tself.merge_threshold = merge_threshold\n\t\tself.cutoff = cutoff\n\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tContract input with list of MPS cores and return result as contractable\n\n\t\tMergedLinearRegion keeps an input counter of the number of inputs, and\n\t\twhen this exceeds its merge threshold, triggers an unmerging and\n\t\tremerging of its parameter tensors.\n\n\t\tArgs:\n\t\t\tinput_data (Tensor): Input with shape [batch_size, input_dim,\n\t\t\t\t\t\t\t\t\t\t\t\t   feature_dim]\n\t\t\"\"\"\n\t\t# If we've hit our threshold, flip the merge state of our tensors\n\t\tif self.input_counter >= self.merge_threshold:\n\t\t\tbond_list, sv_list = self.unmerge(cutoff=self.cutoff)\n\t\t\tself.offset = (self.offset + 1) % 2\n\t\t\tself.merge(offset=self.offset)\n\t\t\tself.input_counter -= self.merge_threshold\n\n\t\t\t# Point self.module_list to the appropriate merged module\n\t\t\tself.module_list = getattr(self, f\"module_list_{self.offset}\")\n\t\telse:\n\t\t\tbond_list, sv_list = None, None\n\n\t\t# Increment our counter and call the LinearRegion's forward method\n\t\tself.input_counter += input_data.size(0)\n\t\toutput = super().forward(input_data)\n\n\t\t# If we flipped our merge state, then return the bond_list and output\n\t\tif bond_list:\n\t\t\treturn output, bond_list, sv_list\n\t\telse:\n\t\t\treturn output\n\n\t@torch.no_grad()\n\tdef merge(self, offset):\n\t\t\"\"\"\n\t\tConvert unmerged modules in self.module_list to merged counterparts\n\n\t\tThis proceeds by first merging all unmerged cores internally, then\n\t\tmerging lone cores when possible during a second sweep\n\t\t\"\"\"\n\t\tassert offset in [0, 1]\n\n\t\tunmerged_list = self.module_list\n\n\t\t# Merge each core internally and add the results to midway_list\n\t\tsite_num = offset\n\t\tmerged_list = []\n\t\tfor core in unmerged_list:\n\t\t\tassert not isinstance(core, MergedInput)\n\t\t\tassert not isinstance(core, MergedOutput)\n\n\t\t\t# Apply internal merging routine if our core supports it\n\t\t\tif hasattr(core, 'merge'):\n\t\t\t\tmerged_list.extend(core.merge(offset=site_num%2))\n\t\t\telse:\n\t\t\t\tmerged_list.append(core)\n\n\t\t\tsite_num += core.core_len()\n\n\t\t# Merge pairs of cores when possible (currently only with\n\t\t# InputSites), making sure to respect the offset for merging.\n\t\twhile True:\n\t\t\tmod_num, site_num = 0, 0\n\t\t\tcombined_list = []\n\n\t\t\twhile mod_num < len(merged_list) - 1:\n\t\t\t\tleft_core, right_core = merged_list[mod_num: mod_num+2]\n\t\t\t\tnew_core = self.combine(left_core, right_core,\n\t\t\t\t\t\t\t\t\t\t\t\t   merging=True)\n\n\t\t\t\t# If cores aren't combinable, move our sliding window by 1\n\t\t\t\tif new_core is None or offset != site_num % 2:\n\t\t\t\t\tcombined_list.append(left_core)\n\t\t\t\t\tmod_num += 1\n\t\t\t\t\tsite_num += left_core.core_len()\n\n\t\t\t\t# If we get something new, move to the next distinct pair\n\t\t\t\telse:\n\t\t\t\t\tassert new_core.core_len() == left_core.core_len() + \\\n\t\t\t\t\t\t\t\t\t\t\t\t  right_core.core_len()\n\t\t\t\t\tcombined_list.append(new_core)\n\t\t\t\t\tmod_num += 2\n\t\t\t\t\tsite_num += new_core.core_len()\n\n\t\t\t\t# Add the last core if there's nothing to merge it with\n\t\t\t\tif mod_num == len(merged_list)-1:\n\t\t\t\t\tcombined_list.append(merged_list[mod_num])\n\t\t\t\t\tmod_num += 1\n\n\t\t\t# We're finished when unmerged_list remains unchanged\n\t\t\tif len(combined_list) == len(merged_list):\n\t\t\t\tbreak\n\t\t\telse:\n\t\t\t\tmerged_list = combined_list\n\n\t\t# Finally, update the appropriate merged module list\n\t\tlist_name = f\"module_list_{offset}\"\n\t\t# If the merged module list hasn't been set yet, initialize it\n\t\tif not hasattr(self, list_name):\n\t\t\tsetattr(self, list_name, nn.ModuleList(merged_list))\n\n\t\t# Otherwise, do an in-place update so that all tensors remain\n\t\t# properly registered with whatever optimizer we use\n\t\telse:\n\t\t\tmodule_list = getattr(self, list_name)\n\t\t\tassert len(module_list) == len(merged_list)\n\t\t\tfor i in range(len(module_list)):\n\t\t\t\tassert module_list[i].tensor.shape == \\\n\t\t\t\t\t   merged_list[i].tensor.shape\n\t\t\t\tmodule_list[i].tensor[:] = merged_list[i].tensor\n\n\t@torch.no_grad()\n\tdef unmerge(self, cutoff=1e-10):\n\t\t\"\"\"\n\t\tConvert merged modules to unmerged counterparts\n\n\t\tThis proceeds by first unmerging all merged cores internally, then\n\t\tcombining lone cores where possible\n\t\t\"\"\"\n\t\tlist_name = f\"module_list_{self.offset}\"\n\t\tmerged_list = getattr(self, list_name)\n\n\t\t# Unmerge each core internally and add results to unmerged_list\n\t\tunmerged_list, bond_list, sv_list = [], [-1], [-1]\n\t\tfor core in merged_list:\n\n\t\t\t# Apply internal unmerging routine if our core supports it\n\t\t\tif hasattr(core, 'unmerge'):\n\t\t\t\tnew_cores, new_bonds, new_svs = core.unmerge(cutoff)\n\t\t\t\tunmerged_list.extend(new_cores)\n\t\t\t\tbond_list.extend(new_bonds[1:])\n\t\t\t\tsv_list.extend(new_svs[1:])\n\t\t\telse:\n\t\t\t\tassert not isinstance(core, InputRegion)\n\t\t\t\tunmerged_list.append(core)\n\t\t\t\tbond_list.append(-1)\n\t\t\t\tsv_list.append(-1)\n\n\t\t# Combine all combinable pairs of cores. This occurs in several\n\t\t# passes, and for now acts nontrivially only on InputSite instances\n\t\twhile True:\n\t\t\tmod_num = 0\n\t\t\tcombined_list = []\n\n\t\t\twhile mod_num < len(unmerged_list) - 1:\n\t\t\t\tleft_core, right_core = unmerged_list[mod_num: mod_num+2]\n\t\t\t\tnew_core = self.combine(left_core, right_core,\n\t\t\t\t\t\t\t\t\t\t\t\t   merging=False)\n\n\t\t\t\t# If cores aren't combinable, move our sliding window by 1\n\t\t\t\tif new_core is None:\n\t\t\t\t\tcombined_list.append(left_core)\n\t\t\t\t\tmod_num += 1\n\n\t\t\t\t# If we get something new, move to the next distinct pair\n\t\t\t\telse:\n\t\t\t\t\tcombined_list.append(new_core)\n\t\t\t\t\tmod_num += 2\n\n\t\t\t\t# Add the last core if there's nothing to combine it with\n\t\t\t\tif mod_num == len(unmerged_list)-1:\n\t\t\t\t\tcombined_list.append(unmerged_list[mod_num])\n\t\t\t\t\tmod_num += 1\n\n\t\t\t# We're finished when unmerged_list remains unchanged\n\t\t\tif len(combined_list) == len(unmerged_list):\n\t\t\t\tbreak\n\t\t\telse:\n\t\t\t\tunmerged_list = combined_list\n\n\t\t# Find the average (log) norm of all of our cores\n\t\tlog_norms = []\n\t\tfor core in unmerged_list:\n\t\t\tlog_norms.append([torch.log(norm) for norm in core.get_norm()])\n\t\tlog_scale = sum([sum(ns) for ns in log_norms])\n\t\tlog_scale /= sum([len(ns) for ns in log_norms])\n\n\t\t# Now rescale all cores so that their norms are roughly equal\n\t\tscales = [[torch.exp(log_scale-n) for n in ns] for ns in log_norms]\n\t\tfor core, these_scales in zip(unmerged_list, scales):\n\t\t\tcore.rescale_norm(these_scales)\n\n\t\t# Add our unmerged module list as a new attribute and return\n\t\t# the updated bond dimensions\n\t\tself.module_list = nn.ModuleList(unmerged_list)\n\t\treturn bond_list, sv_list\n\n\tdef combine(self, left_core, right_core, merging):\n\t\t\"\"\"\n\t\tCombine a pair of cores into a new core using context-dependent rules\n\n\t\tDepending on the types of left_core and right_core, along with whether\n\t\twe're currently merging (merging=True) or unmerging (merging=False),\n\t\teither return a new core, or None if no rule exists for this context\n\t\t\"\"\"\n\n\t\t# Combine an OutputSite with a stray InputSite, return a MergedOutput\n\t\tif merging and ((isinstance(left_core, OutputSite) and\n\t\t\t\t\t\t isinstance(right_core, InputSite)) or\n\t\t\t\t\t\t\t(isinstance(left_core, InputSite) and\n\t\t\t\t\t\t\tisinstance(right_core, OutputSite))):\n\n\t\t\tleft_site = isinstance(left_core, InputSite)\n\t\t\tif left_site:\n\t\t\t\tnew_tensor = torch.einsum('lui,our->olri', [left_core.tensor,\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tright_core.tensor])\n\t\t\telse:\n\t\t\t\tnew_tensor = torch.einsum('olu,uri->olri', [left_core.tensor,\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\tright_core.tensor])\n\t\t\treturn MergedOutput(new_tensor, left_output=(not left_site))\n\n\t\t# Combine an InputRegion with a stray InputSite, return an InputRegion\n\t\telif not merging and ((isinstance(left_core, InputRegion) and\n\t\t\t\t\t\t\t   isinstance(right_core, InputSite)) or\n\t\t\t\t\t\t\t\t\t(isinstance(left_core, InputSite) and\n\t\t\t\t\t\t\t\t\tisinstance(right_core, InputRegion))):\n\n\t\t\tleft_site = isinstance(left_core, InputSite)\n\t\t\tif left_site:\n\t\t\t\tleft_tensor = left_core.tensor.unsqueeze(0)\n\t\t\t\tright_tensor = right_core.tensor\n\t\t\telse:\n\t\t\t\tleft_tensor = left_core.tensor\n\t\t\t\tright_tensor = right_core.tensor.unsqueeze(0)\n\n\t\t\tassert left_tensor.shape[1:] == right_tensor.shape[1:]\n\t\t\tnew_tensor = torch.cat([left_tensor, right_tensor])\n\n\t\t\treturn InputRegion(new_tensor)\n\n\t\t# If this situation doesn't belong to the above cases, return None\n\t\telse:\n\t\t\treturn None\n\n\tdef core_len(self):\n\t\t\"\"\"\n\t\tReturns the number of cores, which is at least the required input size\n\t\t\"\"\"\n\t\treturn sum([module.core_len() for module in self.module_list])\n\n\tdef __len__(self):\n\t\t\"\"\"\n\t\tReturns the number of input sites, which is the required input size\n\t\t\"\"\"\n\t\treturn sum([len(module) for module in self.module_list])\n\nclass InputRegion(nn.Module):\n\t\"\"\"\n\tContiguous region of MPS input cores, associated with bond_str = 'slri'\n\t\"\"\"\n\tdef __init__(self, tensor, use_bias=True, fixed_bias=True, bias_mat=None,\n\t\t\t\t ephemeral=False):\n\t\tsuper().__init__()\n\n\t\t# Make sure tensor has correct size and the component mats are square\n\t\tassert len(tensor.shape) == 4\n\t\tassert tensor.size(1) == tensor.size(2)\n\t\tbond_dim = tensor.size(1)\n\n\t\t# If we are using bias matrices, set those up here\n\t\tif use_bias:\n\t\t\tassert bias_mat is None or isinstance(bias_mat, torch.Tensor)\n\t\t\tbias_mat = torch.eye(bond_dim).unsqueeze(0) if bias_mat is None \\\n\t\t\t\t\t   else bias_mat\n\n\t\t\tbias_modes = len(list(bias_mat.shape))\n\t\t\tassert bias_modes in [2, 3]\n\t\t\tif bias_modes == 2:\n\t\t\t\tbias_mat = bias_mat.unsqueeze(0)\n\n\t\t# Register our tensors as a Pytorch Parameter or Tensor\n\t\tif ephemeral:\n\t\t\tself.register_buffer(name='tensor', tensor=tensor.contiguous())\n\t\t\tself.register_buffer(name='bias_mat', tensor=bias_mat)\n\t\telse:\n\t\t\tself.register_parameter(name='tensor', \n\t\t\t\t\t\t\t\t\tparam=nn.Parameter(tensor.contiguous()))\n\t\t\tif fixed_bias:\n\t\t\t\tself.register_buffer(name='bias_mat', tensor=bias_mat)\n\t\t\telse:\n\t\t\t\tself.register_parameter(name='bias_mat', \n\t\t\t\t\t\t\t\t\t\tparam=nn.Parameter(bias_mat))\n\n\t\tself.use_bias = use_bias\n\t\tself.fixed_bias = fixed_bias\n\t\t\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tContract input with MPS cores and return result as a MatRegion\n\n\t\tArgs:\n\t\t\tinput_data (Tensor): Input with shape [batch_size, input_dim,\n\t\t\t\t\t\t\t\t\t\t\t\t   feature_dim]\n\t\t\"\"\"\n\t\t# Check that input_data has the correct shape\n\t\ttensor = self.tensor\n\t\tassert len(input_data.shape) == 3\n\t\tassert input_data.size(1) == len(self)\n\t\tassert input_data.size(2) == tensor.size(3)\n\n\t\t# Contract the input with our core tensor\n\t\tmats = torch.einsum('slri,bsi->bslr', [tensor, input_data])\n\n\t\t# If we're using bias matrices, add those here\n\t\tif self.use_bias:\n\t\t\tbond_dim = tensor.size(1)\n\t\t\tbias_mat = self.bias_mat.unsqueeze(0)\n\t\t\tmats = mats + bias_mat.expand_as(mats)\n\n\t\treturn MatRegion(mats)\n\n\tdef merge(self, offset):\n\t\t\"\"\"\n\t\tMerge all pairs of neighboring cores and return a new list of cores\n\n\t\toffset is either 0 or 1, which gives the first core at which we start\n\t\tour merging. Depending on the length of our InputRegion, the output of\n\t\tmerge may have 1, 2, or 3 entries, with the majority of sites ending in\n\t\ta MergedInput instance\n\t\t\"\"\"\n\t\tassert offset in [0, 1]\n\t\tnum_sites = self.core_len()\n\t\tparity = num_sites % 2\n\n\t\t# Cases with empty tensors might arise in recursion below\n\t\tif num_sites == 0:\n\t\t\treturn [None]\n\n\t\t# Simplify the problem into one where offset=0 and num_sites is even\n\t\tif (offset, parity) == (1, 1):\n\t\t\tout_list = [self[0], self[1:].merge(offset=0)[0]]\n\t\telif (offset, parity) == (1, 0):\n\t\t\tout_list = [self[0], self[1:-1].merge(offset=0)[0], self[-1]]\n\t\telif (offset, parity) == (0, 1):\n\t\t\tout_list = [self[:-1].merge(offset=0)[0], self[-1]]\n\n\t\t# The main case of interest, with no offset and an even number of sites\n\t\telse:\n\t\t\ttensor = self.tensor\n\t\t\teven_cores, odd_cores = tensor[0::2], tensor[1::2]\n\t\t\tassert len(even_cores) == len(odd_cores)\n\n\t\t\t# Multiply all pairs of cores, keeping inputs separate\n\t\t\tmerged_cores = torch.einsum('slui,surj->slrij', [even_cores,\n\t\t\t\t\t\t\t\t\t\t\t\t\t\t\t odd_cores])\n\t\t\tout_list = [MergedInput(merged_cores)]\n\n\t\t# Remove empty MergedInputs, which appear in very small InputRegions\n\t\treturn [x for x in out_list if x is not None]\n\n\tdef __getitem__(self, key):\n\t\t\"\"\"\n\t\tReturns an InputRegion instance sliced along the site index\n\t\t\"\"\"\n\t\tassert isinstance(key, int) or isinstance(key, slice)\n\n\t\tif isinstance(key, slice):\n\t\t\treturn InputRegion(self.tensor[key])\n\t\telse:\n\t\t\treturn InputSite(self.tensor[key])\n\n\tdef get_norm(self):\n\t\t\"\"\"\n\t\tReturns list of the norms of each core in InputRegion\n\t\t\"\"\"\n\t\treturn [torch.norm(core) for core in self.tensor]\n\n\t@torch.no_grad()\n\tdef rescale_norm(self, scale_list):\n\t\t\"\"\"\n\t\tRescales the norm of each core by an amount specified in scale_list\n\n\t\tFor the i'th tensor defining a core in InputRegion, we rescale as\n\t\ttensor_i <- scale_i * tensor_i, where scale_i = scale_list[i]\n\t\t\"\"\"\n\t\tassert len(scale_list) == len(self.tensor)\n\n\t\tfor core, scale in zip(self.tensor, scale_list):\n\t\t\tcore *= scale\n\n\tdef core_len(self):\n\t\treturn len(self)\n\n\tdef __len__(self):\n\t\treturn self.tensor.size(0)\n\nclass MergedInput(nn.Module):\n\t\"\"\"\n\tContiguous region of merged MPS cores, each taking in a pair of input data\n\n\tSince MergedInput arises after contracting together existing input cores,\n\ta merged input tensor is required for initialization\n\t\"\"\"\n\tdef __init__(self, tensor):\n\t\t# Check that our input tensor has the correct shape\n\t\tbond_str = 'slrij'\n\t\tshape = tensor.shape\n\t\tassert len(shape) == 5\n\t\tassert shape[1] == shape[2]\n\t\tassert shape[3] == shape[4]\n\n\t\tsuper().__init__()\n\n\t\t# Register our tensor as a Pytorch Parameter\n\t\tself.register_parameter(name='tensor', \n\t\t\t\t\t\t\t\tparam=nn.Parameter(tensor.contiguous()))\n\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tContract input with merged MPS cores and return result as a MatRegion\n\n\t\tArgs:\n\t\t\tinput_data (Tensor): Input with shape [batch_size, input_dim,\n\t\t\t\t\t\t\t\t feature_dim], where input_dim must be even\n\t\t\t\t\t\t\t\t (each merged core takes 2 inputs)\n\t\t\"\"\"\n\t\t# Check that input_data has the correct shape\n\t\ttensor = self.tensor\n\t\tassert len(input_data.shape) == 3\n\t\tassert input_data.size(1) == len(self)\n\t\tassert input_data.size(2) == tensor.size(3)\n\t\tassert input_data.size(1) % 2 == 0\n\n\t\t# Divide input_data into inputs living on even and on odd sites\n\t\tinputs = [input_data[:, 0::2], input_data[:, 1::2]]\n\n\t\t# Contract the odd (right-most) and even inputs with merged cores\n\t\ttensor = torch.einsum('slrij,bsj->bslri', [tensor, inputs[1]])\n\t\tmats = torch.einsum('bslri,bsi->bslr', [tensor, inputs[0]])\n\n\t\treturn MatRegion(mats)\n\n\tdef unmerge(self, cutoff=1e-10):\n\t\t\"\"\"\n\t\tSeparate the cores in our MergedInput and return an InputRegion\n\n\t\tThe length of the resultant InputRegion will be identical to our\n\t\toriginal MergedInput (same number of inputs), but its core_len will\n\t\tbe doubled (twice as many individual cores)\n\t\t\"\"\"\n\t\tbond_str = 'slrij'\n\t\ttensor = self.tensor\n\t\tsvd_string = 'lrij->lui,urj'\n\t\tmax_D = tensor.size(1)\n\n\t\t# Split every one of the cores into two and add them both to core_list\n\t\tcore_list, bond_list, sv_list = [], [-1], [-1]\n\t\tfor merged_core in tensor:\n\t\t\tsv_vec = torch.empty(max_D)\n\t\t\tleft_core, right_core, bond_dim = svd_flex(merged_core, svd_string,\n\t\t\t\t\t\t\t\t\t\t\t  max_D, cutoff, sv_vec=sv_vec)\n\n\t\t\tcore_list += [left_core, right_core]\n\t\t\tbond_list += [bond_dim, -1]\n\t\t\tsv_list += [sv_vec, -1]\n\n\t\t# Collate the split cores into one tensor and return as an InputRegion\n\t\ttensor = torch.stack(core_list)\n\t\treturn [InputRegion(tensor)], bond_list, sv_list\n\n\tdef get_norm(self):\n\t\t\"\"\"\n\t\tReturns list of the norm of each core in MergedInput\n\t\t\"\"\"\n\t\treturn [torch.norm(core) for core in self.tensor]\n\n\t@torch.no_grad()\n\tdef rescale_norm(self, scale_list):\n\t\t\"\"\"\n\t\tRescales the norm of each core by an amount specified in scale_list\n\n\t\tFor the i'th tensor defining a core in MergedInput, we rescale as\n\t\ttensor_i <- scale_i * tensor_i, where scale_i = scale_list[i]\n\t\t\"\"\"\n\t\tassert len(scale_list) == len(self.tensor)\n\n\t\tfor core, scale in zip(self.tensor, scale_list):\n\t\t\tcore *= scale\n\n\tdef core_len(self):\n\t\treturn len(self)\n\n\tdef __len__(self):\n\t\t\"\"\"\n\t\tReturns the number of input sites, which is twice the number of cores\n\t\t\"\"\"\n\t\treturn 2 * self.tensor.size(0)\n\nclass InputSite(nn.Module):\n\t\"\"\"\n\tA single MPS core which takes in a single input datum, bond_str = 'lri'\n\t\"\"\"\n\tdef __init__(self, tensor):\n\t\tsuper().__init__()\n\t\t# Register our tensor as a Pytorch Parameter\n\t\tself.register_parameter(name='tensor', \n\t\t\t\t\t\t\t\tparam=nn.Parameter(tensor.contiguous()))\n\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tContract input with MPS core and return result as a SingleMat\n\n\t\tArgs:\n\t\t\tinput_data (Tensor): Input with shape [batch_size, feature_dim]\n\t\t\"\"\"\n\t\t# Check that input_data has the correct shape\n\t\ttensor = self.tensor\n\t\tassert len(input_data.shape) == 2\n\t\tassert input_data.size(1) == tensor.size(2)\n\n\t\t# Contract the input with our core tensor\n\t\tmat = torch.einsum('lri,bi->blr', [tensor, input_data])\n\n\t\treturn SingleMat(mat)\n\n\tdef get_norm(self):\n\t\t\"\"\"\n\t\tReturns the norm of our core tensor, wrapped as a singleton list\n\t\t\"\"\"\n\t\treturn [torch.norm(self.tensor)]\n\n\t@torch.no_grad()\n\tdef rescale_norm(self, scale):\n\t\t\"\"\"\n\t\tRescales the norm of our core by a factor of input `scale`\n\t\t\"\"\"\n\t\tif isinstance(scale, list):\n\t\t\tassert len(scale) == 1\n\t\t\tscale = scale[0]\n\n\t\tself.tensor *= scale\n\n\tdef core_len(self):\n\t\treturn 1\n\n\tdef __len__(self):\n\t\treturn 1\n\nclass OutputSite(nn.Module):\n\t\"\"\"\n\tA single MPS core with no input and a single output index, bond_str = 'olr'\n\t\"\"\"\n\tdef __init__(self, tensor):\n\t\tsuper().__init__()\n\t\t# Register our tensor as a Pytorch Parameter\n\t\tself.register_parameter(name='tensor', \n\t\t\t\t\t\t\t\tparam=nn.Parameter(tensor.contiguous()))\n\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tReturn the OutputSite wrapped as an OutputCore contractable\n\t\t\"\"\"\n\t\treturn OutputCore(self.tensor)\n\n\tdef get_norm(self):\n\t\t\"\"\"\n\t\tReturns the norm of our core tensor, wrapped as a singleton list\n\t\t\"\"\"\n\t\treturn [torch.norm(self.tensor)]\n\n\t@torch.no_grad()\n\tdef rescale_norm(self, scale):\n\t\t\"\"\"\n\t\tRescales the norm of our core by a factor of input `scale`\n\t\t\"\"\"\n\t\tif isinstance(scale, list):\n\t\t\tassert len(scale) == 1\n\t\t\tscale = scale[0]\n\n\t\tself.tensor *= scale\n\n\tdef core_len(self):\n\t\treturn 1\n\n\tdef __len__(self):\n\t\treturn 0\n\nclass MergedOutput(nn.Module):\n\t\"\"\"\n\tMerged MPS core taking in one input datum and returning an output vector\n\n\tSince MergedOutput arises after contracting together an existing input and\n\toutput core, an already-merged tensor is required for initialization\n\n\tArgs:\n\t\ttensor (Tensor):\tValue that our merged core is initialized to\n\t\tleft_output (bool): Specifies if the output core is on the left side of\n\t\t\t\t\t\t\tthe input core (True), or on the right (False)\n\t\"\"\"\n\tdef __init__(self, tensor, left_output):\n\t\t# Check that our input tensor has the correct shape\n\t\tbond_str = 'olri'\n\t\tassert len(tensor.shape) == 4\n\t\tsuper().__init__()\n\n\t\t# Register our tensor as a Pytorch Parameter\n\t\tself.register_parameter(name='tensor', \n\t\t\t\t\t\t\t\tparam=nn.Parameter(tensor.contiguous()))\n\t\tself.left_output = left_output\n\n\tdef forward(self, input_data):\n\t\t\"\"\"\n\t\tContract input with input index of core and return an OutputCore\n\n\t\tArgs:\n\t\t\tinput_data (Tensor): Input with shape [batch_size, feature_dim]\n\t\t\"\"\"\n\t\t# Check that input_data has the correct shape\n\t\ttensor = self.tensor\n\t\tassert len(input_data.shape) == 2\n\t\tassert input_data.size(1) == tensor.size(3)\n\n\t\t# Contract the input with our core tensor\n\t\ttensor = torch.einsum('olri,bi->bolr', [tensor, input_data])\n\n\t\treturn OutputCore(tensor)\n\n\tdef unmerge(self, cutoff=1e-10):\n\t\t\"\"\"\n\t\tSplit our MergedOutput into an OutputSite and an InputSite\n\n\t\tThe non-zero entries of our tensors are dynamically sized according to\n\t\tthe SVD cutoff, but will generally be padded with zeros to give the\n\t\tnew index a regular size.\n\t\t\"\"\"\n\t\tbond_str = 'olri'\n\t\ttensor = self.tensor\n\t\tleft_output = self.left_output\n\t\tif left_output:\n\t\t\tsvd_string = 'olri->olu,uri'\n\t\t\tmax_D = tensor.size(2)\n\t\t\tsv_vec = torch.empty(max_D)\n\n\t\t\toutput_core, input_core, bond_dim = svd_flex(tensor, svd_string,\n\t\t\t\t\t\t\t\t\t\t\t\tmax_D, cutoff, sv_vec=sv_vec)\n\t\t\treturn ([OutputSite(output_core), InputSite(input_core)],\n\t\t\t\t\t[-1, bond_dim, -1], [-1, sv_vec, -1])\n\n\t\telse:\n\t\t\tsvd_string = 'olri->our,lui'\n\t\t\tmax_D = tensor.size(1)\n\t\t\tsv_vec = torch.empty(max_D)\n\n\t\t\toutput_core, input_core, bond_dim = svd_flex(tensor, svd_string,\n\t\t\t\t\t\t\t\t\t\t\t\tmax_D, cutoff, sv_vec=sv_vec)\n\t\t\treturn ([InputSite(input_core), OutputSite(output_core)],\n\t\t\t\t\t[-1, bond_dim, -1], [-1, sv_vec, -1])\n\n\tdef get_norm(self):\n\t\t\"\"\"\n\t\tReturns the norm of our core tensor, wrapped as a singleton list\n\t\t\"\"\"\n\t\treturn [torch.norm(self.tensor)]\n\n\t@torch.no_grad()\n\tdef rescale_norm(self, scale):\n\t\t\"\"\"\n\t\tRescales the norm of our core by a factor of input `scale`\n\t\t\"\"\"\n\t\tif isinstance(scale, list):\n\t\t\tassert len(scale) == 1\n\t\t\tscale = scale[0]\n\n\t\tself.tensor *= scale\n\n\tdef core_len(self):\n\t\treturn 2\n\n\tdef __len__(self):\n\t\treturn 1\n\nclass InitialVector(nn.Module):\n\t\"\"\"\n\tVector of ones and zeros to act as initial vector within the MPS\n\n\tBy default the initial vector is chosen to be all ones, but if fill_dim is\n\tspecified then only the first fill_dim entries are set to one, with the\n\trest zero.\n\n\tIf fixed_vec is False, then the initial vector will be registered as a \n\ttrainable model parameter.\n\t\"\"\"\n\tdef __init__(self, bond_dim, fill_dim=None, fixed_vec=True, \n\t\t\t\t is_left_vec=True):\n\t\tsuper().__init__()\n\n\t\tvec = torch.ones(bond_dim)\n\t\tif fill_dim is not None:\n\t\t\tassert fill_dim >= 0 and fill_dim <= bond_dim\n\t\t\tvec[fill_dim:] = 0\n\n\t\tif fixed_vec:\n\t\t\tvec.requires_grad = False\n\t\t\tself.register_buffer(name='vec', tensor=vec)\n\t\telse:\n\t\t\tvec.requires_grad = True\n\t\t\tself.register_parameter(name='vec', param=nn.Parameter(vec))\n\t\t\n\t\tassert isinstance(is_left_vec, bool)\n\t\tself.is_left_vec = is_left_vec\n\n\tdef forward(self):\n\t\t\"\"\"\n\t\tReturn our initial vector wrapped as an EdgeVec contractable\n\t\t\"\"\"\n\t\treturn EdgeVec(self.vec, self.is_left_vec)\n\n\tdef core_len(self):\n\t\treturn 1\n\n\tdef __len__(self):\n\t\treturn 0\n\nclass TerminalOutput(nn.Module):\n\t\"\"\"\n\tOutput matrix at end of chain to transmute virtual state into output vector\n\n\tBy default, a fixed rectangular identity matrix with shape \n\t[bond_dim, output_dim] will be used as a state transducer. If fixed_mat is\n\tFalse, then the matrix will be registered as a trainable model parameter. \n\t\"\"\"\n\tdef __init__(self, bond_dim, output_dim, fixed_mat=False,\n\t\t\t\t is_left_mat=False):\n\t\tsuper().__init__()\n\n\t\t# I don't have a nice initialization scheme for a non-injective fixed\n\t\t# state transducer, so just throw an error if that's needed\n\t\tif fixed_mat and output_dim > bond_dim:\n\t\t\traise ValueError(\"With fixed_mat=True, TerminalOutput currently \"\n\t\t\t\t\t\t\t \"only supports initialization for bond_dim >= \"\n\t\t\t\t\t\t\t \"output_dim, but here bond_dim=\"\n\t\t\t\t\t\t\tf\"{bond_dim} and output_dim={output_dim}\")\n\n\t\t# Initialize the matrix and register it appropriately\n\t\tmat = torch.eye(bond_dim, output_dim)\n\t\tif fixed_mat:\n\t\t\tmat.requires_grad = False\n\t\t\tself.register_buffer(name='mat', tensor=mat)\n\t\telse:\n\t\t\t# Add some noise to help with training\n\t\t\tmat = mat + torch.randn_like(mat) / bond_dim\n\n\t\t\tmat.requires_grad = True\n\t\t\tself.register_parameter(name='mat', param=nn.Parameter(mat))\n\n\t\tassert isinstance(is_left_mat, bool)\n\t\tself.is_left_mat = is_left_mat\n\n\tdef forward(self):\n\t\t\"\"\"\n\t\tReturn our terminal matrix wrapped as an OutputMat contractable\n\t\t\"\"\"\n\t\treturn OutputMat(self.mat, self.is_left_mat)\n\n\tdef core_len(self):\n\t\treturn 1\n\n\tdef __len__(self):\n\t\treturn 0","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.580207Z","iopub.execute_input":"2021-06-12T19:41:19.580926Z","iopub.status.idle":"2021-06-12T19:41:19.963715Z","shell.execute_reply.started":"2021-06-12T19:41:19.580871Z","shell.execute_reply":"2021-06-12T19:41:19.962912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"EPS = 1e-6\nclass loTeNet(nn.Module):\n    def __init__(self, input_dim, output_dim, bond_dim, feature_dim=2, nCh=3,\n                 kernel=2, virtual_dim=1,\n                 adaptive_mode=False, periodic_bc=False, parallel_eval=False,\n                 label_site=None, path=None, init_std=1e-9, use_bias=True,\n                 fixed_bias=True, cutoff=1e-10, merge_threshold=2000):\n        super().__init__()\n        self.input_dim = input_dim\n        self.virtual_dim = bond_dim\n        \n        ### Squeezing of spatial dimension in first step\t\n        self.kScale = 4\n        nCh =  self.kScale**2 * nCh\n        self.input_dim = self.input_dim // self.kScale\n\n        self.nCh = nCh\n        self.ker = kernel\t\t \n        iDim = (self.input_dim // (self.ker))\n\n        feature_dim = 2*nCh \n        \n        ### First level MPS blocks\n        self.module1 = nn.ModuleList([ MPS(input_dim=(self.ker)**2,\n            output_dim=self.virtual_dim, \n            nCh=nCh, bond_dim=bond_dim, \n            feature_dim=feature_dim, parallel_eval=parallel_eval,\n            adaptive_mode=adaptive_mode, periodic_bc=periodic_bc) \n            for i in range(torch.prod(iDim))])\n\n        self.BN1 = nn.BatchNorm1d(self.virtual_dim,affine=True)\n\n\n        iDim = iDim // self.ker\n        feature_dim = 2*self.virtual_dim\n\n        ### Second level MPS blocks\n        self.module2 = nn.ModuleList([ MPS(input_dim=self.ker**2, \n            output_dim=self.virtual_dim, \n            nCh=self.virtual_dim, bond_dim=bond_dim,\n            feature_dim=feature_dim,  parallel_eval=parallel_eval,\n            adaptive_mode=adaptive_mode, periodic_bc=periodic_bc)\n            for i in range(torch.prod(iDim))])\n\n        self.BN2 = nn.BatchNorm1d(self.virtual_dim,affine=True)\n\n        iDim = iDim // self.ker\n\n        ### Third level MPS blocks\n        self.module3 = nn.ModuleList([ MPS(input_dim=self.ker**2,\n                                           output_dim=self.virtual_dim, \n                                           nCh=self.virtual_dim, bond_dim=bond_dim,  \n                                           feature_dim=feature_dim, parallel_eval=parallel_eval,\n                                           adaptive_mode=adaptive_mode, periodic_bc=periodic_bc) \n        for i in range(torch.prod(iDim))])\n        self.BN3 = nn.BatchNorm1d(self.virtual_dim,affine=True)\n\n        ### Final MPS block\n        self.mpsFinal = MPS(input_dim=len(self.module3), \n                            output_dim=output_dim, nCh=1,\n                            bond_dim=bond_dim, feature_dim=feature_dim, \n                            adaptive_mode=adaptive_mode, periodic_bc=periodic_bc, \n                            parallel_eval=parallel_eval)\n\n    def forward(self,x):\n        b = x.shape[0] #Batch size\n\n        # Increase input feature channel\n        iDim = self.input_dim\n        if self.kScale > 1:\n            x = x.unfold(2,iDim[0],iDim[0]).unfold(3,iDim[1],iDim[1])\n            x = x.reshape(b,iDim[0],iDim[1],-1)\n\n        # Level 1 contraction \n        iDim = self.input_dim//(self.ker)\n        x = x.unfold(2,iDim[0],iDim[0]).unfold(3,iDim[1],iDim[1]).reshape(b,self.nCh,(self.ker)**2,-1)\n        y = [ self.module1[i](x[:,:,:,i]) for i in range(len(self.module1))]\n        y = torch.stack(y,dim=2)\n        y = self.BN1(y)\n\n        # Level 2 contraction\n\n        y = y.view(b,self.virtual_dim,iDim[0],iDim[1])\n        iDim = (iDim//self.ker)\n        y = y.unfold(2,iDim[0],iDim[0]).unfold(3,iDim[1],iDim[1]).reshape(b,self.virtual_dim,self.ker**2,-1)\n        x = [ self.module2[i](y[:,:,:,i]) for i in range(len(self.module2))]\n        x = torch.stack(x,dim=2)\n        x = self.BN2(x)\n\n\n        # Level 3 contraction\n        x = x.view(b,self.virtual_dim,iDim[0],iDim[1])\n        iDim = (iDim//self.ker)\n        x = x.unfold(2,iDim[0],iDim[0]).unfold(3,iDim[1],iDim[1]).reshape(b,self.virtual_dim,self.ker**2,-1)\n        y = [ self.module3[i](x[:,:,:,i]) for i in range(len(self.module3))]\n\n        y = torch.stack(y,dim=2)\n        y = self.BN3(y)\n\n        if y.shape[1] > 1:\n            \n        # Final layer\n            y = self.mpsFinal(y)\n\n        return y.view(b)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:19.969662Z","iopub.execute_input":"2021-06-12T19:41:19.971891Z","iopub.status.idle":"2021-06-12T19:41:20.004358Z","shell.execute_reply.started":"2021-06-12T19:41:19.971824Z","shell.execute_reply":"2021-06-12T19:41:20.003435Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch.optim as optim\n\nimport torch.nn.functional as F\nfrom torch.autograd import Variable\n\nimport torchvision.datasets as dset\nimport torchvision.transforms as transforms\nfrom torch.utils.data import DataLoader\n\nimport torchvision.models as models\n\nimport sys\nimport math\nimport pdb\n\nclass Bottleneck(nn.Module):\n    def __init__(self, nChannels, growthRate):\n        super(Bottleneck, self).__init__()\n        interChannels = 4*growthRate\n        self.bn1 = nn.BatchNorm2d(nChannels)\n        self.conv1 = nn.Conv2d(nChannels, interChannels, kernel_size=1,\n                               bias=False)\n        self.bn2 = nn.BatchNorm2d(interChannels)\n        self.conv2 = nn.Conv2d(interChannels, growthRate, kernel_size=3,\n                               padding=1, bias=False)\n\n    def forward(self, x):\n        out = self.conv1(F.relu(self.bn1(x)))\n        out = self.conv2(F.relu(self.bn2(out)))\n        out = torch.cat((x, out), 1)\n        return out\n\nclass SingleLayer(nn.Module):\n    def __init__(self, nChannels, growthRate):\n        super(SingleLayer, self).__init__()\n        self.bn1 = nn.BatchNorm2d(nChannels)\n        self.conv1 = nn.Conv2d(nChannels, growthRate, kernel_size=3,\n                               padding=1, bias=False)\n\n    def forward(self, x):\n        out = self.conv1(F.relu(self.bn1(x)))\n        out = torch.cat((x, out), 1)\n        return out\n\nclass Transition(nn.Module):\n    def __init__(self, nChannels, nOutChannels):\n        super(Transition, self).__init__()\n        self.bn1 = nn.BatchNorm2d(nChannels)\n        self.conv1 = nn.Conv2d(nChannels, nOutChannels, kernel_size=1,\n                               bias=False)\n\n    def forward(self, x):\n        out = self.conv1(F.relu(self.bn1(x)))\n        out = F.avg_pool2d(out, 2)\n        return out\n\n\nclass DenseNet(nn.Module):\n    def __init__(self, growthRate, depth, reduction, nClasses, bottleneck):\n        super(DenseNet, self).__init__()\n\n        nDenseBlocks = (depth-4) // 3\n        if bottleneck:\n            nDenseBlocks //= 2\n\n        nChannels = 2*growthRate\n        self.conv1 = nn.Conv2d(1, nChannels, kernel_size=3, padding=1,\n                               bias=False)\n        self.dense1 = self._make_dense(nChannels, growthRate, nDenseBlocks, bottleneck)\n        nChannels += nDenseBlocks*growthRate\n        nOutChannels = int(math.floor(nChannels*reduction))\n        self.trans1 = Transition(nChannels, nOutChannels)\n\n        nChannels = nOutChannels\n        self.dense2 = self._make_dense(nChannels, growthRate, nDenseBlocks, bottleneck)\n        nChannels += nDenseBlocks*growthRate\n        nOutChannels = int(math.floor(nChannels*reduction))\n        self.trans2 = Transition(nChannels, nOutChannels)\n\n        nChannels = nOutChannels\n        self.dense3 = self._make_dense(nChannels, growthRate, nDenseBlocks, bottleneck)\n        nChannels += nDenseBlocks*growthRate\n\n        self.bn1 = nn.BatchNorm2d(nChannels)\n        self.fc = nn.Linear(nChannels, nClasses)\n\n        self.nChannels = nChannels\n        for m in self.modules():\n            if isinstance(m, nn.Conv2d):\n                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels\n                m.weight.data.normal_(0, math.sqrt(2. / n))\n            elif isinstance(m, nn.BatchNorm2d):\n                m.weight.data.fill_(1)\n                m.bias.data.zero_()\n            elif isinstance(m, nn.Linear):\n                m.bias.data.zero_()\n\n    def _make_dense(self, nChannels, growthRate, nDenseBlocks, bottleneck):\n        layers = []\n        for i in range(int(nDenseBlocks)):\n            if bottleneck:\n                layers.append(Bottleneck(nChannels, growthRate))\n            else:\n                layers.append(SingleLayer(nChannels, growthRate))\n            nChannels += growthRate\n        return nn.Sequential(*layers)\n\n    def forward(self, x):\n#        pdb.set_trace()\n        out = self.conv1(x)\n        out = self.trans1(self.dense1(out))\n        out = self.trans2(self.dense2(out))\n        out = self.dense3(out)\n        out = torch.squeeze(F.avg_pool2d(F.relu(self.bn1(out)), 8))\n        out = out.view(x.shape[0],self.nChannels,-1).mean(2)\n        out = torch.sigmoid(self.fc(out))\n        return out.squeeze()","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:20.009002Z","iopub.execute_input":"2021-06-12T19:41:20.011387Z","iopub.status.idle":"2021-06-12T19:41:20.045978Z","shell.execute_reply.started":"2021-06-12T19:41:20.011291Z","shell.execute_reply":"2021-06-12T19:41:20.045187Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Initiliaze input dimensions\ndim = torch.ShortTensor(list(train_dataset[0][0].shape[1:]))\nnCh = int(train_dataset[0][0].shape[0])","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:20.050293Z","iopub.execute_input":"2021-06-12T19:41:20.052618Z","iopub.status.idle":"2021-06-12T19:41:20.152565Z","shell.execute_reply.started":"2021-06-12T19:41:20.052574Z","shell.execute_reply":"2021-06-12T19:41:20.151714Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Initialize the models\nif not args.dense_net:\n    print(\"Using LoTeNet\")\n    model = loTeNet(input_dim=dim, output_dim=output_dim, \n                    nCh=nCh, kernel=kernel,\n                    bond_dim=args.bond_dim, feature_dim=feature_dim,\n                    adaptive_mode=adaptive_mode, periodic_bc=periodic_bc, virtual_dim=1)\nelse:\n    print(\"Densenet Baseline!\")\n    model = DenseNet(depth=40, growthRate=12, \n                     reduction=0.5,bottleneck=True,nClasses=output_dim)\nmodel = loTeNet(input_dim=dim, output_dim=output_dim, \n                nCh=nCh, kernel=kernel,\n                bond_dim=args.bond_dim, feature_dim=feature_dim,\n                adaptive_mode=adaptive_mode, periodic_bc=periodic_bc, virtual_dim=1)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:20.156494Z","iopub.execute_input":"2021-06-12T19:41:20.158588Z","iopub.status.idle":"2021-06-12T19:41:58.178840Z","shell.execute_reply.started":"2021-06-12T19:41:20.158536Z","shell.execute_reply":"2021-06-12T19:41:58.177774Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Choose loss function and optimizer\nloss_fun = torch.nn.BCELoss()\noptimizer = torch.optim.Adam(model.parameters(), lr=args.lr, \n                             weight_decay=args.l2)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:58.180542Z","iopub.execute_input":"2021-06-12T19:41:58.180876Z","iopub.status.idle":"2021-06-12T19:41:58.206026Z","shell.execute_reply.started":"2021-06-12T19:41:58.180848Z","shell.execute_reply":"2021-06-12T19:41:58.205108Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"nParam = sum(p.numel() for p in model.parameters() if p.requires_grad)\nprint(\"Number of parameters:%d\"%(nParam))\nprint(f\"Maximum MPS bond dimension = {args.bond_dim}\")\n\nprint(\"Bond dim: %d\"%(args.bond_dim))\nprint(\"Number of parameters:%d\"%(nParam),)","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:58.207219Z","iopub.execute_input":"2021-06-12T19:41:58.207560Z","iopub.status.idle":"2021-06-12T19:41:58.229341Z","shell.execute_reply.started":"2021-06-12T19:41:58.207523Z","shell.execute_reply":"2021-06-12T19:41:58.228523Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(f\"Using Adam w/ learning rate = {args.lr:.1e}\")\nprint(\"Feature_dim: %d, nCh: %d, B:%d\"%(feature_dim,nCh,batch_size))\n\nmodel = model.to(device)\nnValid = len(val_loader)\nnTrain = len(train_loader)\nnTest = len(test_loader)\n\nmaxAuc = 0\nminLoss = 1e3\nconvCheck = 5\nconvIter = 0","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:58.230418Z","iopub.execute_input":"2021-06-12T19:41:58.230908Z","iopub.status.idle":"2021-06-12T19:41:58.266268Z","shell.execute_reply.started":"2021-06-12T19:41:58.230871Z","shell.execute_reply":"2021-06-12T19:41:58.265529Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def evaluate(loader):\n     ### Evaluation funcntion for validation/testing\n\n    with torch.no_grad():\n        vl_acc = 0.\n        vl_loss = 0.\n        labelsNp = np.zeros(1)\n        predsNp = np.zeros(1)\n        model.eval()\n\n        for i, (inputs, labels) in enumerate(loader):\n\n            inputs = inputs.to(device).type(dtype=torch.float)\n            labels = labels.to(device).type(dtype=torch.float)\n            labelsNp = np.concatenate((labelsNp, labels.cpu().numpy()))\n\n            # Inference\n            scores = torch.sigmoid(model(inputs)).type(dtype=torch.float)\n\n            preds = scores\n            loss = loss_fun(scores, labels)\n            predsNp = np.concatenate((predsNp, preds.cpu().numpy()))\n            vl_loss += loss.item()\n\n        # Compute AUC over the full (valid/test) set\n        vl_acc = computeAuc(labelsNp[1:],predsNp[1:])\n        vl_loss = vl_loss/len(loader)\n\n    return vl_acc, vl_loss","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:58.267538Z","iopub.execute_input":"2021-06-12T19:41:58.267881Z","iopub.status.idle":"2021-06-12T19:41:58.275566Z","shell.execute_reply.started":"2021-06-12T19:41:58.267846Z","shell.execute_reply":"2021-06-12T19:41:58.274562Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's start training!\nfor epoch in range(args.num_epochs):\n    running_loss = 0.\n    running_acc = 0.\n    #t = time.time()\n    model.train()\n    predsNp = np.zeros(1)\n    labelsNp = np.zeros(1)\n\n    for i, (inputs, labels) in enumerate(train_loader):\n        \n        # convert inputs and labels  and scores to float tensor\n        inputs = inputs.to(device).type(dtype=torch.float)\n        labels = labels.to(device).type(dtype=torch.float)\n        labelsNp = np.concatenate((labelsNp, labels.cpu().numpy()))\n\n        scores = torch.sigmoid(model(inputs)).type(dtype=torch.float)\n\n        preds = scores\n        loss = loss_fun(scores, labels)\n\n        with torch.no_grad():\n            predsNp = np.concatenate((predsNp, preds.detach().cpu().numpy()))\n            running_loss += loss\n\n        # Backpropagate and update parameters\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        if (i+1) % 5 == 0:\n            print ('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'.format(epoch+1, args.num_epochs, i+1, nTrain, loss.item()))\n\n    accuracy = computeAuc(labelsNp,predsNp)\n\n    # Evaluate on Validation set \n    with torch.no_grad():\n\n        vl_acc, vl_loss = evaluate(val_loader)\n        if vl_acc > maxAuc or vl_loss < minLoss:\n            if vl_loss < minLoss:\n                minLoss = vl_loss\n            if vl_acc > maxAuc:\n                ### Predict on test set\n                ts_acc, ts_loss = evaluate(test_loader)\n                maxAuc = vl_acc\n                print('New Max: %.4f'%maxAuc)\n                print('Test Set Loss:%.4f\tAuc:%.4f'%(ts_loss, ts_acc))\n            convEpoch = epoch\n            convIter = 0\n        else:\n            convIter += 1\n        if convIter == convCheck:\n            if not args.dense_net:\n                print(\"MPS\")\n            else:\n                print(\"DenseNet\")\n            print(\"Converged at epoch:%d with AUC:%.4f\"%(convEpoch+1,maxAuc))\n\n            break","metadata":{"execution":{"iopub.status.busy":"2021-06-12T19:41:58.277248Z","iopub.execute_input":"2021-06-12T19:41:58.277598Z","iopub.status.idle":"2021-06-12T19:42:01.214083Z","shell.execute_reply.started":"2021-06-12T19:41:58.277562Z","shell.execute_reply":"2021-06-12T19:42:01.212810Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}