net.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "import numpy as np\n",
    "import pprint as pp\n",
    "from copy import deepcopy\n",
    "import pickle\n",
    "from numbers import Number\n",
    "from collections import OrderedDict\n",
    "import itertools\n",
    "import pdb\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.autograd import Variable\n",
    "from torch.nn import functional as F\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import ReduceLROnPlateau, LambdaLR, CosineAnnealingLR, CosineAnnealingWarmRestarts\n",
    "from torch.distributions import constraints\n",
    "from torch.distributions.normal import Normal\n",
    "from torch.distributions.multivariate_normal import MultivariateNormal\n",
    "from torch.distributions.distribution import Distribution\n",
    "from torch.distributions.utils import broadcast_all\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(__file__), '..'))\n",
    "from pytorch_net.modules import get_Layer, load_layer_dict, Simple_2_Symbolic\n",
    "from pytorch_net.util import forward, get_epochs_T_mult, Loss_Fun, get_activation, get_criterion, get_criteria_value, get_optimizer, get_full_struct_param, plot_matrices, get_model_DL, PrecisionFloorLoss, get_list_DL, init_weight\n",
    "from pytorch_net.util import Early_Stopping, Performance_Monitor, record_data, to_np_array, to_Variable, make_dir, formalize_value, RampupLR, Transform_Label, view_item, load_model, save_model, to_cpu_recur, filter_kwargs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training functionality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(\n",
    "    model,\n",
    "    X=None,\n",
    "    y=None,\n",
    "    train_loader=None,\n",
    "    validation_data=None,\n",
    "    validation_loader=None,\n",
    "    criterion=nn.MSELoss(),\n",
    "    inspect_interval=10,\n",
    "    isplot=False,\n",
    "    is_cuda=None,\n",
    "    **kwargs\n",
    "    ):\n",
    "    \"\"\"Training function for generic models. \"model\" can be a single model or a ordered list of models\"\"\"\n",
    "    def get_regularization(model, loss_epoch, **kwargs):\n",
    "        \"\"\"Compute regularization.\"\"\"\n",
    "        reg_dict = kwargs[\"reg_dict\"] if \"reg_dict\" in kwargs else None\n",
    "        reg = to_Variable([0], is_cuda = is_cuda)\n",
    "        if reg_dict is not None:\n",
    "            for reg_type, reg_coeff in reg_dict.items():\n",
    "                # Setting up regularization strength:\n",
    "                if isinstance(reg_coeff, Number):\n",
    "                    reg_coeff_ele = reg_coeff\n",
    "                else:\n",
    "                    if loss_epoch < len(reg_coeff):\n",
    "                        reg_coeff_ele = reg_coeff[loss_epoch]\n",
    "                    else:\n",
    "                        reg_coeff_ele = reg_coeff[-1]\n",
    "                # Accumulate regularization:\n",
    "                reg = reg + model.get_regularization(source=[reg_type], mode=reg_mode, **kwargs) * reg_coeff_ele\n",
    "        return reg\n",
    "\n",
    "    if is_cuda is None:\n",
    "        if X is None and y is None:\n",
    "            assert train_loader is not None\n",
    "            is_cuda = train_loader.dataset.tensors[0].is_cuda\n",
    "        else:\n",
    "            is_cuda = X.is_cuda\n",
    "\n",
    "    # Optimization kwargs:\n",
    "    epochs = kwargs[\"epochs\"] if \"epochs\" in kwargs else 10000\n",
    "    lr = kwargs[\"lr\"] if \"lr\" in kwargs else 5e-3\n",
    "    lr_rampup_steps = kwargs[\"lr_rampup\"] if \"lr_rampup\" in kwargs else 200\n",
    "    optim_type = kwargs[\"optim_type\"] if \"optim_type\" in kwargs else \"adam\"\n",
    "    optim_kwargs = kwargs[\"optim_kwargs\"] if \"optim_kwargs\" in kwargs else {}\n",
    "    scheduler_type = kwargs[\"scheduler_type\"] if \"scheduler_type\" in kwargs else \"ReduceLROnPlateau\"\n",
    "    gradient_noise = kwargs[\"gradient_noise\"] if \"gradient_noise\" in kwargs else None\n",
    "    data_loader_apply = kwargs[\"data_loader_apply\"] if \"data_loader_apply\" in kwargs else None\n",
    "\n",
    "    # Inspection kwargs:\n",
    "    inspect_step = kwargs[\"inspect_step\"] if \"inspect_step\" in kwargs else None  # Whether to inspect each step\n",
    "    inspect_items = kwargs[\"inspect_items\"] if \"inspect_items\" in kwargs else None\n",
    "    inspect_items_train = get_inspect_items_train(inspect_items)\n",
    "    inspect_functions = kwargs[\"inspect_functions\"] if \"inspect_functions\" in kwargs else None\n",
    "    if inspect_functions is not None:\n",
    "        for inspect_function_key in inspect_functions:\n",
    "            if inspect_function_key not in inspect_items:\n",
    "                inspect_items.append(inspect_function_key)\n",
    "    inspect_items_interval = kwargs[\"inspect_items_interval\"] if \"inspect_items_interval\" in kwargs else 1000\n",
    "    inspect_image_interval = kwargs[\"inspect_image_interval\"] if \"inspect_image_interval\" in kwargs else None\n",
    "    inspect_loss_precision = kwargs[\"inspect_loss_precision\"] if \"inspect_loss_precision\" in kwargs else 4\n",
    "    callback = kwargs[\"callback\"] if \"callback\" in kwargs else None\n",
    "\n",
    "    # Saving kwargs:\n",
    "    record_keys = kwargs[\"record_keys\"] if \"record_keys\" in kwargs else [\"loss\"]\n",
    "    filename = kwargs[\"filename\"] if \"filename\" in kwargs else None\n",
    "    if filename is not None:\n",
    "        make_dir(filename)\n",
    "    save_interval = kwargs[\"save_interval\"] if \"save_interval\" in kwargs else None\n",
    "    save_step = kwargs[\"save_step\"] if \"save_step\" in kwargs else None\n",
    "    logdir = kwargs[\"logdir\"] if \"logdir\" in kwargs else None\n",
    "    data_record = {key: [] for key in record_keys}\n",
    "    info_to_save = kwargs[\"info_to_save\"] if \"info_to_save\" in kwargs else None\n",
    "    if info_to_save is not None:\n",
    "        data_record.update(info_to_save)\n",
    "    patience = kwargs[\"patience\"] if \"patience\" in kwargs else 20\n",
    "    if patience is not None:\n",
    "        early_stopping_epsilon = kwargs[\"early_stopping_epsilon\"] if \"early_stopping_epsilon\" in kwargs else 0\n",
    "        early_stopping_monitor = kwargs[\"early_stopping_monitor\"] if \"early_stopping_monitor\" in kwargs else \"loss\"\n",
    "        early_stopping = Early_Stopping(patience = patience, epsilon = early_stopping_epsilon, mode = \"max\" if early_stopping_monitor in [\"accuracy\"] else \"min\")\n",
    "    if logdir is not None:\n",
    "        from pytorch_net.logger import Logger\n",
    "        batch_idx = 0\n",
    "        logger = Logger(logdir)\n",
    "    logimages = kwargs[\"logimages\"] if \"logimages\" in kwargs else None\n",
    "    reg_mode = kwargs[\"reg_mode\"] if \"reg_mode\" in kwargs else \"L1\"\n",
    "\n",
    "    if validation_loader is not None:\n",
    "        assert validation_data is None\n",
    "        X_valid, y_valid = None, None\n",
    "    elif validation_data is not None:\n",
    "        X_valid, y_valid = validation_data\n",
    "    else:\n",
    "        X_valid, y_valid = X, y\n",
    "\n",
    "    # Setting up dynamic label noise:\n",
    "    label_noise_matrix = kwargs[\"label_noise_matrix\"] if \"label_noise_matrix\" in kwargs else None\n",
    "    transform_label = Transform_Label(label_noise_matrix = label_noise_matrix, is_cuda=is_cuda)\n",
    "\n",
    "    # Setting up cotrain optimizer:\n",
    "    co_kwargs = kwargs[\"co_kwargs\"] if \"co_kwargs\" in kwargs else None\n",
    "    if co_kwargs is not None:\n",
    "        co_optimizer = co_kwargs[\"co_optimizer\"]\n",
    "        co_model = co_kwargs[\"co_model\"]\n",
    "        co_criterion = co_kwargs[\"co_criterion\"] if \"co_criterion\" in co_kwargs else None\n",
    "        co_multi_step = co_kwargs[\"co_multi_step\"] if \"co_multi_step\" in co_kwargs else 1\n",
    "\n",
    "    # Get original loss:\n",
    "    if len(inspect_items_train) > 0:\n",
    "        loss_value_train = get_loss(model, train_loader, X, y, criterion=criterion, loss_epoch=-1, transform_label=transform_label, **kwargs)\n",
    "        info_dict_train = prepare_inspection(model, train_loader, X, y, transform_label=transform_label, **kwargs)\n",
    "        if \"loss\" in record_keys:\n",
    "            record_data(data_record, [loss_value_train], [\"loss_tr\"])\n",
    "    loss_original = get_loss(model, validation_loader, X_valid, y_valid, criterion=criterion, loss_epoch=-1, transform_label=transform_label, **kwargs)\n",
    "    if \"loss\" in record_keys:\n",
    "        record_data(data_record, [-1, loss_original], [\"iter\", \"loss\"])\n",
    "    if \"reg\" in record_keys and \"reg_dict\" in kwargs and len(kwargs[\"reg_dict\"]) > 0:\n",
    "        reg_value = get_regularization(model, loss_epoch=0, **kwargs)\n",
    "        record_data(data_record, [reg_value], [\"reg\"])\n",
    "    if \"param\" in record_keys:\n",
    "        record_data(data_record, [model.get_weights_bias(W_source=\"core\", b_source=\"core\")], [\"param\"])\n",
    "    if \"param_grad\" in record_keys:\n",
    "        record_data(data_record, [model.get_weights_bias(W_source=\"core\", b_source=\"core\", is_grad=True)], [\"param_grad\"])\n",
    "    if co_kwargs is not None:\n",
    "        co_loss_original = get_loss(co_model, validation_loader, X_valid, y_valid, criterion=criterion, loss_epoch=-1, transform_label=transform_label, **co_kwargs)\n",
    "        if \"co_loss\" in record_keys:\n",
    "            record_data(data_record, [co_loss_original], [\"co_loss\"])\n",
    "    if filename is not None and save_interval is not None:\n",
    "        record_data(data_record, [{}], [\"model_dict\"])\n",
    "\n",
    "    # Setting up optimizer:\n",
    "    parameters = model.parameters()\n",
    "    num_params = len(list(model.parameters()))\n",
    "    if num_params == 0:\n",
    "        print(\"No parameters to optimize!\")\n",
    "        loss_value = get_loss(model, validation_loader, X_valid, y_valid, criterion = criterion, loss_epoch = -1, transform_label=transform_label, **kwargs)\n",
    "        if \"loss\" in record_keys:\n",
    "            record_data(data_record, [0, loss_value], [\"iter\", \"loss\"])\n",
    "        if \"param\" in record_keys:\n",
    "            record_data(data_record, [model.get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "        if \"param_grad\" in record_keys:\n",
    "            record_data(data_record, [model.get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True)], [\"param_grad\"])\n",
    "        if co_kwargs is not None:\n",
    "            co_loss_value = get_loss(co_model, validation_loader, X_valid, y_valid, criterion = criterion, loss_epoch = -1, transform_label=transform_label, **co_kwargs)\n",
    "            record_data(data_record, [co_loss_value], [\"co_loss\"])\n",
    "        return loss_original, loss_value, data_record\n",
    "    optimizer = get_optimizer(optim_type, lr, parameters, **optim_kwargs) if \"optimizer\" not in kwargs or (\"optimizer\" in kwargs and kwargs[\"optimizer\"] is None) else kwargs[\"optimizer\"]\n",
    "\n",
    "    # Initialize inspect_items:\n",
    "    if inspect_items is not None:\n",
    "        print(\"{}:\".format(-1), end = \"\")\n",
    "        print(\"\\tlr: {0:.3e}\\t loss:{1:.{2}f}\".format(optimizer.param_groups[0][\"lr\"], loss_original, inspect_loss_precision), end = \"\")\n",
    "        info_dict = prepare_inspection(model, validation_loader, X_valid, y_valid, transform_label=transform_label, **kwargs)\n",
    "        if len(inspect_items_train) > 0:\n",
    "            print(\"\\tloss_tr: {0:.{1}f}\".format(loss_value_train, inspect_loss_precision), end = \"\")\n",
    "            info_dict_train = update_key_train(info_dict_train, inspect_items_train)\n",
    "            info_dict.update(info_dict_train)\n",
    "        if \"reg\" in record_keys and \"reg_dict\" in kwargs and len(kwargs[\"reg_dict\"]) > 0:\n",
    "            print(\"\\treg:{0:.{1}f}\".format(to_np_array(reg_value), inspect_loss_precision), end=\"\")\n",
    "        if len(info_dict) > 0:\n",
    "            for item in inspect_items:\n",
    "                if item in info_dict:\n",
    "                    print(\" \\t{0}: {1:.{2}f}\".format(item, info_dict[item], inspect_loss_precision), end = \"\")\n",
    "                    if item in record_keys and item not in [\"loss\", \"reg\"]:\n",
    "                        record_data(data_record, [to_np_array(info_dict[item])], [item])\n",
    "\n",
    "        if co_kwargs is not None:\n",
    "            co_info_dict = prepare_inspection(co_model, validation_loader, X_valid, y_valid, transform_label=transform_label, **co_kwargs)\n",
    "            if \"co_loss\" in inspect_items:\n",
    "                co_loss_value = get_loss(co_model, validation_loader, X_valid, y_valid, criterion=criterion, loss_epoch=-1, transform_label=transform_label, **co_kwargs)\n",
    "                print(\"\\tco_loss: {}\".format(formalize_value(co_loss_value, inspect_loss_precision)), end=\"\")\n",
    "            if len(co_info_dict) > 0:\n",
    "                for item in inspect_items:\n",
    "                    if item in co_info_dict:\n",
    "                        print(\" \\t{0}: {1}\".format(item, formalize_value(co_info_dict[item], inspect_loss_precision)), end=\"\")\n",
    "                        if item in record_keys and item != \"loss\":\n",
    "                            record_data(data_record, [to_np_array(co_info_dict[item])], [item])\n",
    "        print(\"\\n\")\n",
    "\n",
    "    # Setting up gradient noise:\n",
    "    if gradient_noise is not None:\n",
    "        from pytorch_net.util import Gradient_Noise_Scale_Gen\n",
    "        scale_gen = Gradient_Noise_Scale_Gen(epochs=epochs,\n",
    "                                             gamma=gradient_noise[\"gamma\"],  # decay rate\n",
    "                                             eta=gradient_noise[\"eta\"],      # starting variance\n",
    "                                             gradient_noise_interval_epoch=1,\n",
    "                                            )\n",
    "        gradient_noise_scale = scale_gen.generate_scale(verbose=True)\n",
    "\n",
    "    # Set up learning rate scheduler:\n",
    "    if scheduler_type is not None:\n",
    "        if scheduler_type == \"ReduceLROnPlateau\":\n",
    "            scheduler_patience = kwargs[\"scheduler_patience\"] if \"scheduler_patience\" in kwargs else 40\n",
    "            scheduler_factor = kwargs[\"scheduler_factor\"] if \"scheduler_factor\" in kwargs else 0.1\n",
    "            scheduler_verbose = kwargs[\"scheduler_verbose\"] if \"scheduler_verbose\" in kwargs else False\n",
    "            scheduler = ReduceLROnPlateau(optimizer, factor=scheduler_factor, patience=scheduler_patience, verbose=scheduler_verbose)\n",
    "        elif scheduler_type == \"LambdaLR\":\n",
    "            scheduler_lr_lambda = kwargs[\"scheduler_lr_lambda\"] if \"scheduler_lr_lambda\" in kwargs else (lambda epoch: 0.97 ** (epoch // 2))\n",
    "            scheduler = LambdaLR(optimizer, lr_lambda=scheduler_lr_lambda)\n",
    "        elif scheduler_type == \"cos\":\n",
    "            scheduler = CosineAnnealingLR(optimizer, T_max=epochs)\n",
    "        elif scheduler_type == \"coslr\":\n",
    "            T_0 = max(min(25, epochs//31), 1)\n",
    "            T_mult = kwargs[\"scheduler_T_mult\"] if \"scheduler_T_mult\" in kwargs else 2\n",
    "            epochs = get_epochs_T_mul(epochs, T_0=T_0, T_mult=T_mult)\n",
    "            scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=T_0, T_mult=T_mult)\n",
    "        else:\n",
    "            raise\n",
    "    # Ramping or learning rate for the first lr_rampup_steps steps:\n",
    "    if lr_rampup_steps is not None and train_loader is not None:\n",
    "        scheduler_rampup = RampupLR(optimizer, num_steps=lr_rampup_steps)\n",
    "        if hasattr(train_loader, \"dataset\"):\n",
    "            data_size = len(train_loader.dataset)\n",
    "        else:\n",
    "            data_size = kwargs[\"data_size\"]\n",
    "\n",
    "    # Initialize logdir:\n",
    "    if logdir is not None:\n",
    "        if logimages is not None:\n",
    "            for tag, image_fun in logimages[\"image_fun\"].items():\n",
    "                image = image_fun(model, logimages[\"X\"], logimages[\"y\"])\n",
    "                logger.log_images(tag, image, -1)\n",
    "\n",
    "    # Training:\n",
    "    to_stop = False\n",
    "\n",
    "    for i in range(epochs + 1):\n",
    "        model.train()\n",
    "\n",
    "        # Updating gradient noise:\n",
    "        if gradient_noise is not None:\n",
    "            hook_handle_list = []\n",
    "            if i % scale_gen.gradient_noise_interval_epoch == 0:\n",
    "                for h in hook_handle_list:\n",
    "                    h.remove()\n",
    "                hook_handle_list = []\n",
    "                scale_idx = int(i / scale_gen.gradient_noise_interval_epoch)\n",
    "                if scale_idx >= len(gradient_noise_scale):\n",
    "                    current_gradient_noise_scale = gradient_noise_scale[-1]\n",
    "                else:\n",
    "                    current_gradient_noise_scale = gradient_noise_scale[scale_idx]\n",
    "                for param_group in optimizer.param_groups:\n",
    "                    for param in param_group[\"params\"]:\n",
    "                        if param.requires_grad:\n",
    "                            h = param.register_hook(lambda grad: grad + Variable(torch.normal(mean=torch.zeros(grad.size()),\n",
    "                                                                                              std=current_gradient_noise_scale * torch.ones(grad.size()))))\n",
    "                            hook_handle_list.append(h)\n",
    "\n",
    "        if X is not None and y is not None:\n",
    "            if optim_type != \"LBFGS\":\n",
    "                optimizer.zero_grad()\n",
    "                reg = get_regularization(model, loss_epoch=i, **kwargs)\n",
    "                loss = model.get_loss(X, transform_label(y), criterion=criterion, loss_epoch=i, **kwargs) + reg\n",
    "                loss.backward()\n",
    "                optimizer.step()\n",
    "            else:\n",
    "                # \"LBFGS\" is a second-order optimization algorithm that requires a slightly different procedure:\n",
    "                def closure():\n",
    "                    optimizer.zero_grad()\n",
    "                    reg = get_regularization(model, loss_epoch=i, **kwargs)\n",
    "                    loss = model.get_loss(X, transform_label(y), criterion=criterion, loss_epoch=i, **kwargs) + reg\n",
    "                    loss.backward()\n",
    "                    return loss\n",
    "                optimizer.step(closure)\n",
    "            \n",
    "            # Cotrain step:\n",
    "            if co_kwargs is not None:\n",
    "                if \"co_warmup_epochs\" not in co_kwargs or \"co_warmup_epochs\" in co_kwargs and i >= co_kwargs[\"co_warmup_epochs\"]:\n",
    "                    for _ in range(co_multi_step):\n",
    "                        co_optimizer.zero_grad()\n",
    "                        co_reg = get_regularization(co_model, loss_epoch=i, **co_kwargs)\n",
    "                        co_loss = co_model.get_loss(X, transform_label(y), criterion=co_criterion, loss_epoch=i, **co_kwargs) + co_reg\n",
    "                        co_loss.backward()\n",
    "                        co_optimizer.step()\n",
    "        else:\n",
    "            if inspect_step is not None:\n",
    "                info_dict_step = {key: [] for key in inspect_items}\n",
    "\n",
    "            if \"loader_process\" in kwargs and kwargs[\"loader_process\"] is not None:\n",
    "                train_loader = kwargs[\"loader_process\"](\"train\")\n",
    "            for k, data_batch in enumerate(train_loader):\n",
    "                if isinstance(data_batch, tuple) or isinstance(data_batch, list):\n",
    "                    X_batch, y_batch = data_batch\n",
    "                    if data_loader_apply is not None:\n",
    "                        X_batch, y_batch = data_loader_apply(X_batch, y_batch)\n",
    "                else:\n",
    "                    X_batch, y_batch = data_loader_apply(data_batch)\n",
    "                if optim_type != \"LBFGS\":\n",
    "                    optimizer.zero_grad()\n",
    "                    reg = get_regularization(model, loss_epoch=i, **kwargs)\n",
    "                    loss = model.get_loss(X_batch, transform_label(y_batch), criterion=criterion, loss_epoch=i, loss_step=k, **kwargs) + reg\n",
    "                    loss.backward()\n",
    "                    if logdir is not None:\n",
    "                        batch_idx += 1\n",
    "                        if len(model.info_dict) > 0:\n",
    "                            for item in inspect_items:\n",
    "                                if item in model.info_dict:\n",
    "                                    logger.log_scalar(item, model.info_dict[item], batch_idx)\n",
    "                    optimizer.step()\n",
    "                else:\n",
    "                    def closure():\n",
    "                        optimizer.zero_grad()\n",
    "                        reg = get_regularization(model, loss_epoch=i, **kwargs)\n",
    "                        loss = model.get_loss(X_batch, transform_label(y_batch), criterion=criterion, loss_epoch=i, loss_step=k, **kwargs) + reg\n",
    "                        loss.backward()\n",
    "                        return loss\n",
    "                    if logdir is not None:\n",
    "                        batch_idx += 1\n",
    "                        if len(model.info_dict) > 0:\n",
    "                            for item in inspect_items:\n",
    "                                if item in model.info_dict:\n",
    "                                    logger.log_scalar(item, model.info_dict[item], batch_idx)\n",
    "                    optimizer.step(closure)\n",
    "\n",
    "                # Rampup scheduler:\n",
    "                if lr_rampup_steps is not None and i * data_size // len(X_batch) + k < lr_rampup_steps:\n",
    "                    scheduler_rampup.step()\n",
    "\n",
    "                # Cotrain step:\n",
    "                if co_kwargs is not None:\n",
    "                    if \"co_warmup_epochs\" not in co_kwargs or \"co_warmup_epochs\" in co_kwargs and i >= co_kwargs[\"co_warmup_epochs\"]:\n",
    "                        for _ in range(co_multi_step):\n",
    "                            co_optimizer.zero_grad()\n",
    "                            co_reg = get_regularization(co_model, loss_epoch=i, **co_kwargs)\n",
    "                            co_loss = co_model.get_loss(X_batch, transform_label(y_batch), criterion=co_criterion, loss_epoch=i, loss_step=k, **co_kwargs) + co_reg\n",
    "                            co_loss.backward()\n",
    "                            if logdir is not None:\n",
    "                                if len(co_model.info_dict) > 0:\n",
    "                                    for item in inspect_items:\n",
    "                                        if item in co_model.info_dict:\n",
    "                                            logger.log_scalar(item, co_model.info_dict[item], batch_idx)\n",
    "                            co_optimizer.step()\n",
    "\n",
    "                # Inspect at each step:\n",
    "                if inspect_step is not None:\n",
    "                    if k % inspect_step == 0:\n",
    "                        print(\"s{}:\".format(k), end = \"\")\n",
    "                        info_dict = prepare_inspection(model, validation_loader, X_valid, y_valid, transform_label=transform_label, **kwargs) \n",
    "                        if \"loss\" in inspect_items:\n",
    "                            info_dict_step[\"loss\"].append(loss.item())\n",
    "                            print(\"\\tloss: {0:.{1}f}\".format(loss.item(), inspect_loss_precision), end=\"\")\n",
    "                        if len(info_dict) > 0:\n",
    "                            for item in inspect_items:\n",
    "                                if item in info_dict:\n",
    "                                    info_dict_step[item].append(info_dict[item])\n",
    "                                    print(\" \\t{0}: {1}\".format(item, formalize_value(info_dict[item], inspect_loss_precision)), end = \"\")\n",
    "                        if co_kwargs is not None:\n",
    "                            if \"co_warmup_epochs\" not in co_kwargs or \"co_warmup_epochs\" in co_kwargs and i >= co_kwargs[\"co_warmup_epochs\"]:\n",
    "                                co_info_dict = prepare_inspection(co_model, validation_loader, X_valid, y_valid, transform_label=transform_label, **co_kwargs)\n",
    "                                if \"co_loss\" in inspect_items:\n",
    "                                    print(\"\\tco_loss: {0:.{1}f}\".format(co_loss.item(), inspect_loss_precision), end=\"\")\n",
    "                                    info_dict_step[\"co_loss\"].append(co_loss.item())\n",
    "                                if len(co_info_dict) > 0:\n",
    "                                    for item in inspect_items:\n",
    "                                        if item in co_info_dict and item != \"co_loss\":\n",
    "                                            info_dict_step[item].append(co_info_dict[item])\n",
    "                                            print(\" \\t{0}: {1}\".format(item, formalize_value(co_info_dict[item], inspect_loss_precision)), end=\"\")\n",
    "                        print()\n",
    "                    if k % save_step == 0:\n",
    "                        if filename is not None:\n",
    "                            pickle.dump(model.model_dict, open(filename[:-2] + \"_model.p\", \"wb\"))\n",
    "\n",
    "        if logdir is not None:\n",
    "            # Log values and gradients of the parameters (histogram summary)\n",
    "#             for tag, value in model.named_parameters():\n",
    "#                 tag = tag.replace('.', '/')\n",
    "#                 logger.log_histogram(tag, to_np_array(value), i)\n",
    "#                 logger.log_histogram(tag + '/grad', to_np_array(value.grad), i)\n",
    "            if logimages is not None:\n",
    "                for tag, image_fun in logimages[\"image_fun\"].items():\n",
    "                    image = image_fun(model, logimages[\"X\"], logimages[\"y\"])\n",
    "                    logger.log_images(tag, image, i)\n",
    "\n",
    "        if i % inspect_interval == 0:\n",
    "            model.eval()\n",
    "            if inspect_items is not None and i % inspect_items_interval == 0 and len(inspect_items_train) > 0:\n",
    "                loss_value_train = get_loss(model, train_loader, X, y, criterion = criterion, loss_epoch = i, transform_label=transform_label, **kwargs)\n",
    "                info_dict_train = prepare_inspection(model, train_loader, X, y, transform_label=transform_label, **kwargs)\n",
    "            loss_value = get_loss(model, validation_loader, X_valid, y_valid, criterion = criterion, loss_epoch = i, transform_label=transform_label, **kwargs)\n",
    "            reg_value = get_regularization(model, loss_epoch = i, **kwargs)\n",
    "            if scheduler_type is not None:\n",
    "                if lr_rampup_steps is None or train_loader is None or (lr_rampup_steps is not None and i * data_size // len(X_batch) + k >= lr_rampup_steps):\n",
    "                    if scheduler_type == \"ReduceLROnPlateau\":\n",
    "                        scheduler.step(loss_value)\n",
    "                    else:\n",
    "                        scheduler.step()\n",
    "            if callback is not None:\n",
    "                assert callable(callback)\n",
    "                callback(model = model,\n",
    "                         X = X_valid,\n",
    "                         y = y_valid,\n",
    "                         iteration = i,\n",
    "                         loss = loss_value,\n",
    "                        )\n",
    "            if patience is not None:\n",
    "                if early_stopping_monitor == \"loss\":\n",
    "                    to_stop = early_stopping.monitor(loss_value)\n",
    "                else:\n",
    "                    info_dict = prepare_inspection(model, validation_loader, X_valid, y_valid, transform_label=transform_label, **kwargs)\n",
    "                    to_stop = early_stopping.monitor(info_dict[early_stopping_monitor])\n",
    "            if inspect_items is not None:\n",
    "                if i % inspect_items_interval == 0:\n",
    "                    # Get loss:\n",
    "                    print(\"{}:\".format(i), end = \"\")\n",
    "                    print(\"\\tlr: {0:.3e}\\tloss: {1:.{2}f}\".format(optimizer.param_groups[0][\"lr\"], loss_value, inspect_loss_precision), end = \"\")\n",
    "                    info_dict = prepare_inspection(model, validation_loader, X_valid, y_valid, transform_label=transform_label, **kwargs)\n",
    "                    if len(inspect_items_train) > 0:\n",
    "                        print(\"\\tloss_tr: {0:.{1}f}\".format(loss_value_train, inspect_loss_precision), end = \"\")\n",
    "                        info_dict_train = update_key_train(info_dict_train, inspect_items_train)\n",
    "                        info_dict.update(info_dict_train)\n",
    "                    if \"reg\" in inspect_items and \"reg_dict\" in kwargs and len(kwargs[\"reg_dict\"]) > 0:\n",
    "                        print(\"\\treg:{0:.{1}f}\".format(to_np_array(reg_value), inspect_loss_precision), end=\"\")\n",
    "                    \n",
    "                    # Print and record:\n",
    "                    if len(info_dict) > 0:\n",
    "                        for item in inspect_items:\n",
    "                            if item + \"_val\" in info_dict:\n",
    "                                print(\" \\t{0}: {1}\".format(item, formalize_value(info_dict[item + \"_val\"], inspect_loss_precision)), end = \"\")\n",
    "                                if item in record_keys and item not in [\"loss\", \"reg\"]:\n",
    "                                    record_data(data_record, [to_np_array(info_dict[item + \"_val\"])], [item])\n",
    "\n",
    "                        # logger:\n",
    "                        if logdir is not None:\n",
    "                            for item in inspect_items:\n",
    "                                if item + \"_val\" in info_dict:\n",
    "                                    logger.log_scalar(item + \"_val\", info_dict[item + \"_val\"], i)\n",
    "\n",
    "                    # Co_model:\n",
    "                    if co_kwargs is not None:\n",
    "                        co_loss_value = get_loss(co_model, validation_loader, X_valid, y_valid, criterion = criterion, loss_epoch = i, transform_label=transform_label, **co_kwargs)\n",
    "                        co_info_dict = prepare_inspection(co_model, validation_loader, X_valid, y_valid, transform_label=transform_label, **co_kwargs)\n",
    "                        if \"co_loss\" in inspect_items:\n",
    "                            print(\"\\tco_loss: {0:.{1}f}\".format(co_loss_value, inspect_loss_precision), end=\"\")\n",
    "                        if len(co_info_dict) > 0:\n",
    "                            for item in inspect_items:\n",
    "                                if item + \"_val\" in co_info_dict:\n",
    "                                    print(\" \\t{0}: {1}\".format(item, formalize_value(co_info_dict[item + \"_val\"], inspect_loss_precision)), end=\"\")\n",
    "                                    if item in record_keys and item != \"co_loss\":\n",
    "                                        record_data(data_record, [to_np_array(co_info_dict[item + \"_val\"])], [item])\n",
    "                        if \"co_loss\" in record_keys:\n",
    "                            record_data(data_record, [co_loss_value], [\"co_loss\"])\n",
    "\n",
    "                    # Training metrics:\n",
    "                    if inspect_step is not None:\n",
    "                        for item in info_dict_step:\n",
    "                            if len(info_dict_step[item]) > 0:\n",
    "                                print(\" \\t{0}_s: {1}\".format(item, formalize_value(np.mean(info_dict_step[item]), inspect_loss_precision)), end = \"\")\n",
    "                                if item in record_keys and item != \"loss\":\n",
    "                                    record_data(data_record, [np.mean(info_dict_step[item])], [\"{}_s\".format(item)])\n",
    "\n",
    "                    # Record loss:\n",
    "                    if \"loss\" in record_keys:\n",
    "                        record_data(data_record, [i, loss_value], [\"iter\", \"loss\"])\n",
    "                    if \"reg\" in record_keys and \"reg_dict\" in kwargs and len(kwargs[\"reg_dict\"]) > 0:\n",
    "                        record_data(data_record, [reg_value], [\"reg\"])\n",
    "                    if \"param\" in record_keys:\n",
    "                        record_data(data_record, [model.get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "                    if \"param_grad\" in record_keys:\n",
    "                        record_data(data_record, [model.get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True)], [\"param_grad\"])\n",
    "                    print(\"\\n\")\n",
    "                    try:\n",
    "                        sys.stdout.flush()\n",
    "                    except:\n",
    "                        pass\n",
    "            if isplot:\n",
    "                if inspect_image_interval is not None and hasattr(model, \"plot\"):\n",
    "                    if i % inspect_image_interval == 0:\n",
    "                        if gradient_noise is not None:\n",
    "                            print(\"gradient_noise: {0:.9f}\".format(current_gradient_noise_scale))\n",
    "                        plot_model(model, data_loader = validation_loader, X = X_valid, y = y_valid, transform_label=transform_label, data_loader_apply=data_loader_apply)\n",
    "                if co_kwargs is not None and \"inspect_image_interval\" in co_kwargs and co_kwargs[\"inspect_image_interval\"] and hasattr(co_model, \"plot\"):\n",
    "                    if i % co_kwargs[\"inspect_image_interval\"] == 0:\n",
    "                        plot_model(co_model, data_loader = validation_loader, X = X_valid, y = y_valid, transform_label=transform_label, data_loader_apply=data_loader_apply)\n",
    "        if save_interval is not None:\n",
    "            if i % save_interval == 0:\n",
    "                record_data(data_record, [model.model_dict], [\"model_dict\"])\n",
    "                if co_kwargs is not None:\n",
    "                    record_data(data_record, [co_model.model_dict], [\"co_model_dict\"])\n",
    "                if filename is not None:\n",
    "                    pickle.dump(data_record, open(filename, \"wb\"))\n",
    "        if to_stop:\n",
    "            break\n",
    "\n",
    "    loss_value = get_loss(model, validation_loader, X_valid, y_valid, criterion=criterion, loss_epoch=epochs, transform_label=transform_label, **kwargs)\n",
    "    if isplot:\n",
    "        import matplotlib.pylab as plt\n",
    "        for key, item in data_record.items():\n",
    "            if isinstance(item, Number) or len(data_record[\"iter\"]) != len(item):\n",
    "                continue\n",
    "            if key not in [\"iter\", \"model_dict\"]:\n",
    "                if key in [\"accuracy\"]:\n",
    "                    plt.figure(figsize = (8,6))\n",
    "                    plt.plot(data_record[\"iter\"], data_record[key])\n",
    "                    plt.xlabel(\"epoch\")\n",
    "                    plt.ylabel(key)\n",
    "                    plt.title(key)\n",
    "                    plt.show()\n",
    "                else:\n",
    "                    plt.figure(figsize = (8,6))\n",
    "                    plt.semilogy(data_record[\"iter\"], data_record[key])\n",
    "                    plt.xlabel(\"epoch\")\n",
    "                    plt.ylabel(key)\n",
    "                    plt.title(key)\n",
    "                    plt.show()\n",
    "    return loss_original, loss_value, data_record\n",
    "\n",
    "\n",
    "def train_simple(model, X, y, validation_data = None, inspect_interval = 5, **kwargs):\n",
    "    \"\"\"minimal version of training. \"model\" can be a single model or a ordered list of models\"\"\"\n",
    "    def get_regularization(model, **kwargs):\n",
    "        reg_dict = kwargs[\"reg_dict\"] if \"reg_dict\" in kwargs else None\n",
    "        reg = to_Variable([0], is_cuda = X.is_cuda)\n",
    "        for model_ele in model:\n",
    "            if reg_dict is not None:\n",
    "                for reg_type, reg_coeff in reg_dict.items():\n",
    "                    reg = reg + model_ele.get_regularization(source = [reg_type], mode = \"L1\", **kwargs) * reg_coeff\n",
    "        return reg\n",
    "    if not(isinstance(model, list) or isinstance(model, tuple)):\n",
    "        model = [model]\n",
    "    epochs = kwargs[\"epochs\"] if \"epochs\" in kwargs else 2000\n",
    "    lr = kwargs[\"lr\"] if \"lr\" in kwargs else 5e-3\n",
    "    optim_type = kwargs[\"optim_type\"] if \"optim_type\" in kwargs else \"adam\"\n",
    "    optim_kwargs = kwargs[\"optim_kwargs\"] if \"optim_kwargs\" in kwargs else {}\n",
    "    loss_type = kwargs[\"loss_type\"] if \"loss_type\" in kwargs else \"mse\"\n",
    "    early_stopping_epsilon = kwargs[\"early_stopping_epsilon\"] if \"early_stopping_epsilon\" in kwargs else 0\n",
    "    patience = kwargs[\"patience\"] if \"patience\" in kwargs else 40\n",
    "    record_keys = kwargs[\"record_keys\"] if \"record_keys\" in kwargs else [\"loss\", \"mse\", \"data_DL\", \"model_DL\"]\n",
    "    scheduler_type = kwargs[\"scheduler_type\"] if \"scheduler_type\" in kwargs else \"ReduceLROnPlateau\"\n",
    "    loss_precision_floor = kwargs[\"loss_precision_floor\"] if \"loss_precision_floor\" in kwargs else PrecisionFloorLoss\n",
    "    autoencoder = kwargs[\"autoencoder\"] if \"autoencoder\" in kwargs else None\n",
    "    data_record = {key: [] for key in record_keys}\n",
    "    isplot = kwargs[\"isplot\"] if \"isplot\" in kwargs else False\n",
    "    if patience is not None:\n",
    "        early_stopping = Early_Stopping(patience = patience, epsilon = early_stopping_epsilon)\n",
    "    \n",
    "    if validation_data is not None:\n",
    "        X_valid, y_valid = validation_data\n",
    "    else:\n",
    "        X_valid, y_valid = X, y\n",
    "    \n",
    "    # Get original loss:\n",
    "    criterion = get_criterion(loss_type, loss_precision_floor = loss_precision_floor)\n",
    "    DL_criterion = Loss_Fun(core = \"DLs\", loss_precision_floor = loss_precision_floor, DL_sum = True)\n",
    "    DL_criterion_absolute = Loss_Fun(core = \"DLs\", loss_precision_floor = PrecisionFloorLoss, DL_sum = True)\n",
    "    pred_valid = forward(model, X_valid, **kwargs)\n",
    "    loss_original = to_np_array(criterion(pred_valid, y_valid))\n",
    "    if \"loss\" in record_keys:\n",
    "        record_data(data_record, [-1, loss_original], [\"iter\",\"loss\"])\n",
    "    if \"mse\" in record_keys:\n",
    "        record_data(data_record, [to_np_array(nn.MSELoss()(pred_valid, y_valid))], [\"mse\"])\n",
    "    if \"data_DL\" in record_keys:\n",
    "        record_data(data_record, [to_np_array(DL_criterion(pred_valid, y_valid))], [\"data_DL\"])\n",
    "    if \"data_DL_absolute\" in record_keys:\n",
    "        record_data(data_record, [to_np_array(DL_criterion_absolute(pred_valid, y_valid))], [\"data_DL_absolute\"])\n",
    "    if \"model_DL\" in record_keys:\n",
    "        record_data(data_record, [get_model_DL(model)], [\"model_DL\"])\n",
    "    if \"param\" in record_keys:\n",
    "        record_data(data_record, [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "    if \"param_grad\" in record_keys:\n",
    "        record_data(data_record, [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True)], [\"param_grad\"])\n",
    "    if \"param_collapse_layers\" in record_keys:\n",
    "        record_data(data_record, [simplify(deepcopy(model[0]), X, y, \"collapse_layers\", verbose = 0)[0]\\\n",
    "                                  .get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "\n",
    "    # Setting up optimizer:\n",
    "    parameters = itertools.chain(*[model_ele.parameters() for model_ele in model])\n",
    "    num_params = np.sum([[len(list(model_ele.parameters())) for model_ele in model]])\n",
    "    if num_params == 0:\n",
    "        print(\"No parameters to optimize!\")\n",
    "        pred_valid = forward(model, X_valid, **kwargs)\n",
    "        loss_value = to_np_array(criterion(pred_valid, y_valid))\n",
    "        if \"loss\" in record_keys:\n",
    "            record_data(data_record, [0, loss_value], [\"iter\", \"loss\"])\n",
    "        if \"mse\" in record_keys:\n",
    "            record_data(data_record, [to_np_array(nn.MSELoss()(pred_valid, y_valid))], [\"mse\"])\n",
    "        if \"data_DL\" in record_keys:\n",
    "            record_data(data_record, [to_np_array(DL_criterion(pred_valid, y_valid))], [\"data_DL\"])\n",
    "        if \"data_DL_absolute\" in record_keys:\n",
    "            record_data(data_record, [to_np_array(DL_criterion_absolute(pred_valid, y_valid))], [\"data_DL_absolute\"])\n",
    "        if \"model_DL\" in record_keys:\n",
    "            record_data(data_record, [get_model_DL(model)], [\"model_DL\"])\n",
    "        if \"param\" in record_keys:\n",
    "            record_data(data_record, [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "        if \"param_grad\" in record_keys:\n",
    "            record_data(data_record, [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True)], [\"param_grad\"])\n",
    "        if \"param_collapse_layers\" in record_keys:\n",
    "            record_data(data_record, [simplify(deepcopy(model[0]), X, y, \"collapse_layers\", verbose = 0)[0]\\\n",
    "                                      .get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "        return loss_original, loss_value, data_record\n",
    "    optimizer = get_optimizer(optim_type, lr, parameters, **optim_kwargs)\n",
    "    \n",
    "    # Set up learning rate scheduler:\n",
    "    if scheduler_type is not None:\n",
    "        if scheduler_type == \"ReduceLROnPlateau\":\n",
    "            scheduler_patience = kwargs[\"scheduler_patience\"] if \"scheduler_patience\" in kwargs else 10\n",
    "            scheduler_factor = kwargs[\"scheduler_factor\"] if \"scheduler_factor\" in kwargs else 0.1\n",
    "            scheduler = ReduceLROnPlateau(optimizer, factor = scheduler_factor, patience = scheduler_patience)\n",
    "        elif scheduler_type == \"LambdaLR\":\n",
    "            scheduler_lr_lambda = kwargs[\"scheduler_lr_lambda\"] if \"scheduler_lr_lambda\" in kwargs else (lambda epoch: 1 / (1 + 0.01 * epoch))\n",
    "            scheduler = LambdaLR(optimizer, lr_lambda = scheduler_lr_lambda)\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "    # Training:\n",
    "    to_stop = False\n",
    "    for i in range(epochs + 1):\n",
    "        if optim_type != \"LBFGS\":\n",
    "            optimizer.zero_grad()\n",
    "            pred = forward(model, X, **kwargs)\n",
    "            reg = get_regularization(model, **kwargs)\n",
    "            loss = criterion(pred, y) + reg\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        else:\n",
    "            # \"LBFGS\" is a second-order optimization algorithm that requires a slightly different procedure:\n",
    "            def closure():\n",
    "                optimizer.zero_grad()\n",
    "                pred = forward(model, X, **kwargs)\n",
    "                reg = get_regularization(model, **kwargs)\n",
    "                loss = criterion(pred, y) + reg\n",
    "                loss.backward()\n",
    "                return loss\n",
    "            optimizer.step(closure)\n",
    "        if i % inspect_interval == 0:\n",
    "            pred_valid = forward(model, X_valid, **kwargs)\n",
    "            loss_value = to_np_array(criterion(pred_valid, y_valid))\n",
    "            if scheduler_type is not None:\n",
    "                if scheduler_type == \"ReduceLROnPlateau\":\n",
    "                    scheduler.step(loss_value)\n",
    "                else:\n",
    "                    scheduler.step()\n",
    "            if \"loss\" in record_keys:\n",
    "                record_data(data_record, [i, loss_value], [\"iter\", \"loss\"])\n",
    "            if \"mse\" in record_keys:\n",
    "                record_data(data_record, [to_np_array(nn.MSELoss()(pred_valid, y_valid))], [\"mse\"])\n",
    "            if \"data_DL\" in record_keys:\n",
    "                record_data(data_record, [to_np_array(DL_criterion(pred_valid, y_valid))], [\"data_DL\"])\n",
    "            if \"data_DL_absolute\" in record_keys:\n",
    "                record_data(data_record, [to_np_array(DL_criterion_absolute(pred_valid, y_valid))], [\"data_DL_absolute\"])\n",
    "            if \"model_DL\" in record_keys:\n",
    "                record_data(data_record, [get_model_DL(model)], [\"model_DL\"])\n",
    "            if \"param\" in record_keys:\n",
    "                record_data(data_record, [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "            if \"param_grad\" in record_keys:\n",
    "                record_data(data_record, [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True)], [\"param_grad\"])\n",
    "            if \"param_collapse_layers\" in record_keys:\n",
    "                record_data(data_record, [simplify(deepcopy(model[0]), X, y, \"collapse_layers\", verbose = 0)[0]\\\n",
    "                                          .get_weights_bias(W_source = \"core\", b_source = \"core\")], [\"param\"])\n",
    "            if patience is not None:\n",
    "                to_stop = early_stopping.monitor(loss_value)\n",
    "        if to_stop:\n",
    "            break\n",
    "\n",
    "    pred_valid = forward(model, X_valid, **kwargs)\n",
    "    loss_value = to_np_array(criterion(pred_valid, y_valid))\n",
    "    if isplot:\n",
    "        import matplotlib.pylab as plt\n",
    "        if \"mse\" in data_record:\n",
    "            plt.semilogy(data_record[\"iter\"], data_record[\"mse\"])\n",
    "            plt.xlabel(\"epochs\")\n",
    "            plt.title(\"MSE\")\n",
    "            plt.show()\n",
    "        if \"loss\" in data_record:\n",
    "            plt.plot(data_record[\"iter\"], data_record[\"loss\"])\n",
    "            plt.xlabel(\"epochs\")\n",
    "            plt.title(\"Loss\")\n",
    "            plt.show()\n",
    "    return loss_original, loss_value, data_record\n",
    "\n",
    "\n",
    "def load_model_dict_net(model_dict, is_cuda = False):\n",
    "    net_type = model_dict[\"type\"]\n",
    "    if net_type.startswith(\"MLP\"):\n",
    "        return MLP(input_size = model_dict[\"input_size\"],\n",
    "                   struct_param = model_dict[\"struct_param\"] if \"struct_param\" in model_dict else None,\n",
    "                   W_init_list = model_dict[\"weights\"] if \"weights\" in model_dict else None,\n",
    "                   b_init_list = model_dict[\"bias\"] if \"bias\" in model_dict else None,\n",
    "                   settings = model_dict[\"settings\"] if \"settings\" in model_dict else {},\n",
    "                   is_cuda = is_cuda,\n",
    "                  )\n",
    "    elif net_type == \"Labelmix_MLP\":\n",
    "        model = Labelmix_MLP(input_size=model_dict[\"input_size\"],\n",
    "                             struct_param=model_dict[\"struct_param\"],\n",
    "                             idx_label=model_dict[\"idx_label\"] if \"idx_label\" in model_dict else None,\n",
    "                             is_cuda=is_cuda,\n",
    "                            )\n",
    "        if \"state_dict\" in model_dict:\n",
    "            model.load_state_dict(model_dict[\"state_dict\"])\n",
    "        return model\n",
    "    elif net_type == \"Multi_MLP\":\n",
    "        return Multi_MLP(input_size = model_dict[\"input_size\"],\n",
    "                   struct_param = model_dict[\"struct_param\"],\n",
    "                   W_init_list = model_dict[\"weights\"] if \"weights\" in model_dict else None,\n",
    "                   b_init_list = model_dict[\"bias\"] if \"bias\" in model_dict else None,\n",
    "                   settings = model_dict[\"settings\"] if \"settings\" in model_dict else {},\n",
    "                   is_cuda = is_cuda,\n",
    "                  )\n",
    "    elif net_type == \"Branching_Net\":\n",
    "        return Branching_Net(net_base_model_dict = model_dict[\"net_base_model_dict\"],\n",
    "                             net_1_model_dict = model_dict[\"net_1_model_dict\"],\n",
    "                             net_2_model_dict = model_dict[\"net_2_model_dict\"],\n",
    "                             is_cuda = is_cuda,\n",
    "                            )\n",
    "    elif net_type == \"Fan_in_MLP\":\n",
    "        return Fan_in_MLP(model_dict_branch1=model_dict[\"model_dict_branch1\"],\n",
    "                          model_dict_branch2=model_dict[\"model_dict_branch2\"],\n",
    "                          model_dict_joint=model_dict[\"model_dict_joint\"],\n",
    "                          is_cuda=is_cuda,\n",
    "                         )\n",
    "    elif net_type == \"Net_reparam\":\n",
    "        return Net_reparam(model_dict=model_dict[\"model\"],\n",
    "                           reparam_mode=model_dict[\"reparam_mode\"],\n",
    "                           is_cuda=is_cuda,\n",
    "                          )\n",
    "    elif net_type == \"Wide_ResNet\":\n",
    "        model = Wide_ResNet(depth=model_dict[\"depth\"],\n",
    "                            widen_factor=model_dict[\"widen_factor\"],\n",
    "                            input_channels=model_dict[\"input_channels\"],\n",
    "                            output_size=model_dict[\"output_size\"],\n",
    "                            dropout_rate=model_dict[\"dropout_rate\"],\n",
    "                            is_cuda=is_cuda,\n",
    "                           )\n",
    "        if \"state_dict\" in model_dict:\n",
    "            model.load_state_dict(model_dict[\"state_dict\"])\n",
    "        return model\n",
    "    elif net_type.startswith(\"ConvNet\"):\n",
    "        return ConvNet(input_channels = model_dict[\"input_channels\"],\n",
    "                       struct_param = model_dict[\"struct_param\"],\n",
    "                       W_init_list = model_dict[\"weights\"] if \"weights\" in model_dict else None,\n",
    "                       b_init_list = model_dict[\"bias\"] if \"bias\" in model_dict else None,\n",
    "                       settings = model_dict[\"settings\"] if \"settings\" in model_dict else {},\n",
    "                       return_indices = model_dict[\"return_indices\"] if \"return_indices\" in model_dict else False,\n",
    "                       is_cuda = is_cuda,\n",
    "                      )\n",
    "    elif net_type == \"Conv_Autoencoder\":\n",
    "        model = Conv_Autoencoder(input_channels_encoder = model_dict[\"input_channels_encoder\"],\n",
    "                                 input_channels_decoder = model_dict[\"input_channels_decoder\"],\n",
    "                                 struct_param_encoder = model_dict[\"struct_param_encoder\"],\n",
    "                                 struct_param_decoder = model_dict[\"struct_param_decoder\"],\n",
    "                                 settings = model_dict[\"settings\"],\n",
    "                                 is_cuda = is_cuda,\n",
    "                                )\n",
    "        if \"encoder\" in model_dict:\n",
    "            model.encoder.load_model_dict(model_dict[\"encoder\"])\n",
    "        if \"decoder\" in model_dict:\n",
    "            model.decoder.load_model_dict(model_dict[\"decoder\"])\n",
    "        return model\n",
    "    elif model_dict[\"type\"] == \"Conv_Model\":\n",
    "        is_generative = model_dict[\"is_generative\"] if \"is_generative\" in model_dict else False\n",
    "        return Conv_Model(encoder_model_dict = model_dict[\"encoder_model_dict\"] if not is_generative else None,\n",
    "                          core_model_dict = model_dict[\"core_model_dict\"],\n",
    "                          decoder_model_dict = model_dict[\"decoder_model_dict\"],\n",
    "                          latent_size = model_dict[\"latent_size\"],\n",
    "                          is_generative = model_dict[\"is_generative\"] if is_generative else False,\n",
    "                          is_res_block = model_dict[\"is_res_block\"] if \"is_res_block\" in model_dict else False,\n",
    "                          is_cuda = is_cuda,\n",
    "                         )\n",
    "    else:\n",
    "        raise Exception(\"net_type {} not recognized!\".format(net_type))\n",
    "        \n",
    "\n",
    "def load_model_dict(model_dict, is_cuda = False):\n",
    "    net_type = model_dict[\"type\"]\n",
    "    if net_type not in [\"Model_Ensemble\", \"LSTM\", \"Model_with_Uncertainty\", \"Mixture_Model\", \"Mixture_Gaussian\"]:\n",
    "        return load_model_dict_net(model_dict, is_cuda = is_cuda)\n",
    "    elif net_type == \"Model_Ensemble\":\n",
    "        if model_dict[\"model_type\"] == \"MLP\":\n",
    "            model_ensemble = Model_Ensemble(\n",
    "                num_models = model_dict[\"num_models\"],\n",
    "                input_size = model_dict[\"input_size\"],\n",
    "                model_type = model_dict[\"model_type\"],\n",
    "                output_size = model_dict[\"output_size\"],\n",
    "                is_cuda = is_cuda,\n",
    "                # Here we just create some placeholder network. The model will be overwritten in the next steps:\n",
    "                struct_param = [[1, \"Simple_Layer\", {}]],\n",
    "            )\n",
    "        elif model_dict[\"model_type\"] == \"LSTM\":\n",
    "            model_ensemble = Model_Ensemble(\n",
    "                num_models = model_dict[\"num_models\"],\n",
    "                input_size = model_dict[\"input_size\"],\n",
    "                model_type = model_dict[\"model_type\"],\n",
    "                output_size = model_dict[\"output_size\"],\n",
    "                is_cuda = is_cuda,\n",
    "                # Here we just create some placeholder network. The model will be overwritten in the next steps:\n",
    "                hidden_size = 3,\n",
    "                output_struct_param = [[1, \"Simple_Layer\", {}]],\n",
    "            )\n",
    "        else:\n",
    "            raise\n",
    "        for k in range(model_ensemble.num_models):\n",
    "            setattr(model_ensemble, \"model_{}\".format(k), load_model_dict(model_dict[\"model_{}\".format(k)], is_cuda = is_cuda))\n",
    "        return model_ensemble\n",
    "    elif net_type == \"Model_with_Uncertainty\":\n",
    "        return Model_with_Uncertainty(model_pred = load_model_dict(model_dict[\"model_pred\"], is_cuda = is_cuda),\n",
    "                                      model_logstd = load_model_dict(model_dict[\"model_logstd\"], is_cuda = is_cuda))\n",
    "    elif net_type == \"Mixture_Model\":\n",
    "        return Mixture_Model(model_dict_list=model_dict[\"model_dict_list\"],\n",
    "                             weight_logits_model_dict=model_dict[\"weight_logits_model_dict\"],\n",
    "                             num_components=model_dict[\"num_components\"],\n",
    "                             is_cuda=is_cuda,\n",
    "                            )\n",
    "    elif net_type == \"Mixture_Gaussian\":\n",
    "        return load_model_dict_Mixture_Gaussian(model_dict, is_cuda = is_cuda)\n",
    "    else:\n",
    "        raise Exception(\"net_type {} not recognized!\".format(net_type))\n",
    "\n",
    "\n",
    "## Helper functions:\n",
    "def get_accuracy(pred, target):\n",
    "    \"\"\"Get accuracy from prediction and target\"\"\"\n",
    "    assert len(pred.shape) == len(target.shape) == 1\n",
    "    assert len(pred) == len(target)\n",
    "    pred, target = to_np_array(pred, target)\n",
    "    accuracy = ((pred == target).sum().astype(float) / len(pred))\n",
    "    return accuracy\n",
    "\n",
    "\n",
    "def flatten(*tensors):\n",
    "    \"\"\"Flatten the tensor except the first dimension\"\"\"\n",
    "    new_tensors = []\n",
    "    for tensor in tensors:\n",
    "        new_tensors.append(tensor.view(tensor.size(0), -1))\n",
    "    if len(new_tensors) == 1:\n",
    "        new_tensors = new_tensors[0]\n",
    "    return new_tensors\n",
    "\n",
    "\n",
    "def fill_triangular(vec, dim, mode=\"lower\"):\n",
    "    \"\"\"Fill an lower or upper triangular matrices with given vectors.\n",
    "\n",
    "    Specifically, it transform the examples in the vector form [n_examples, size]\n",
    "        into the lower-triangular matrix with shape [n_examples, dim, dim]\n",
    "    \n",
    "    Args:\n",
    "        vec: [n_examples, size], where the size\n",
    "        dim: the dimension of the lower triangular matrix.\n",
    "            the size == dim * (dim + 1) // 2 must be satisfied.\n",
    "\n",
    "    Returns:\n",
    "        matrix: with shape [n_examples, dim, dim]\n",
    "    \"\"\"\n",
    "#     num_examples, size = vec.shape\n",
    "#     assert size == dim * (dim + 1) // 2\n",
    "#     matrix = torch.zeros(num_examples, dim, dim).to(vec.device)\n",
    "#     if mode == \"lower\":\n",
    "#         idx = (torch.tril(torch.ones(dim, dim)) == 1)[None]\n",
    "#     elif mode == \"upper\":\n",
    "#         idx = (torch.triu(torch.ones(dim, dim)) == 1)[None]\n",
    "#     else:\n",
    "#         raise Exception(\"mode {} not recognized!\".format(mode))\n",
    "#     idx = idx.repeat(num_examples,1,1)\n",
    "#     matrix[idx] = vec.contiguous().view(-1)\n",
    "    num_examples, size = vec.shape\n",
    "    assert size == dim * (dim + 1) // 2\n",
    "    if mode == \"lower\":\n",
    "        rows, cols = torch.tril_indices(dim, dim)\n",
    "    elif mode == \"upper\":\n",
    "        rows, cols = torch.triu_indices(dim, dim)\n",
    "    else:\n",
    "        raise Exception(\"mode {} not recognized!\".format(mode))\n",
    "    matrix = torch.zeros(num_examples, dim, dim).type(vec.dtype).to(vec.device)\n",
    "    matrix[:, rows, cols] = vec\n",
    "    return matrix\n",
    "\n",
    "\n",
    "def matrix_diag_transform(matrix, fun):\n",
    "    \"\"\"Return the matrices whose diagonal elements have been executed by the function 'fun'.\"\"\"\n",
    "    num_examples = len(matrix)\n",
    "    idx = torch.eye(matrix.size(-1)).bool().unsqueeze(0)\n",
    "    idx = idx.repeat(num_examples, 1, 1)\n",
    "    new_matrix = matrix.clone()\n",
    "    new_matrix[idx] = fun(matrix.diagonal(dim1=1, dim2=2).contiguous().view(-1))\n",
    "    return new_matrix\n",
    "\n",
    "\n",
    "def Zip(*data, **kwargs):\n",
    "    \"\"\"Recursive unzipping of data structure\n",
    "    Example: Zip(*[(('a',2), 1), (('b',3), 2), (('c',3), 3), (('d',2), 4)])\n",
    "    ==> [[['a', 'b', 'c', 'd'], [2, 3, 3, 2]], [1, 2, 3, 4]]\n",
    "    Each subtree in the original data must be in the form of a tuple.\n",
    "    In the **kwargs, you can set the function that is applied to each fully unzipped subtree.\n",
    "    \"\"\"\n",
    "    import collections\n",
    "    function = kwargs[\"function\"] if \"function\" in kwargs else None\n",
    "    if len(data) == 1:\n",
    "        return data[0]\n",
    "    data = [list(element) for element in zip(*data)]\n",
    "    for i, element in enumerate(data):\n",
    "        if isinstance(element[0], tuple):\n",
    "            data[i] = Zip(*element, **kwargs)\n",
    "        elif isinstance(element, list):\n",
    "            if function is not None:\n",
    "                data[i] = function(element)\n",
    "    return data\n",
    "\n",
    "\n",
    "def get_loss(model, data_loader=None, X=None, y=None, criterion=None, transform_label=None, **kwargs):\n",
    "    \"\"\"Get loss using the whole data or data_loader. Return the average validation loss with np.ndarray format\"\"\"\n",
    "    max_validation_iter = kwargs[\"max_validation_iter\"] if \"max_validation_iter\" in kwargs else None\n",
    "    if transform_label is None:\n",
    "        transform_label = Transform_Label()\n",
    "    if \"loader_process\" in kwargs and kwargs[\"loader_process\"] is not None:\n",
    "        data_loader = kwargs[\"loader_process\"](\"test\")\n",
    "    if data_loader is not None:\n",
    "        assert X is None and y is None\n",
    "        loss_record = 0\n",
    "        count = 0\n",
    "        # Taking the average of all metrics:\n",
    "        for j, data_batch in enumerate(data_loader):\n",
    "            if isinstance(data_batch, tuple) or isinstance(data_batch, list):\n",
    "                X_batch, y_batch = data_batch\n",
    "                if \"data_loader_apply\" in kwargs and kwargs[\"data_loader_apply\"] is not None:\n",
    "                    X_batch, y_batch = kwargs[\"data_loader_apply\"](X_batch, y_batch)\n",
    "            else:\n",
    "                X_batch, y_batch = kwargs[\"data_loader_apply\"](data_batch)\n",
    "            loss_ele = to_np_array(model.get_loss(X_batch, transform_label(y_batch), criterion = criterion, **kwargs))\n",
    "            if j == 0:\n",
    "                all_info_dict = {key: 0 for key in model.info_dict.keys()}\n",
    "            loss_record = loss_record + loss_ele\n",
    "            count += 1\n",
    "            for key in model.info_dict:\n",
    "                all_info_dict[key] = all_info_dict[key] + model.info_dict[key]\n",
    "\n",
    "            if max_validation_iter is not None and count > max_validation_iter:\n",
    "                break\n",
    "\n",
    "        for key in model.info_dict:\n",
    "            all_info_dict[key] = all_info_dict[key] / count\n",
    "        loss = loss_record / count\n",
    "        model.info_dict = deepcopy(all_info_dict)\n",
    "    else:\n",
    "        assert X is not None and y is not None\n",
    "        loss = to_np_array(model.get_loss(X, transform_label(y), criterion = criterion, **kwargs))\n",
    "    return loss\n",
    "\n",
    "\n",
    "def plot_model(model, data_loader=None, X=None, y=None, transform_label=None, **kwargs):\n",
    "    data_loader_apply = kwargs[\"data_loader_apply\"] if \"data_loader_apply\" in kwargs else None\n",
    "    max_validation_iter = kwargs[\"max_validation_iter\"] if \"max_validation_iter\" in kwargs else None\n",
    "    if transform_label is None:\n",
    "        transform_label = Transform_Label()\n",
    "    if \"loader_process\" in kwargs and kwargs[\"loader_process\"] is not None:\n",
    "        data_loader = kwargs[\"loader_process\"](\"test\")\n",
    "    if data_loader is not None:\n",
    "        assert X is None and y is None\n",
    "        X_all = []\n",
    "        y_all = []\n",
    "        for i, data_batch in enumerate(data_loader):\n",
    "            if isinstance(data_batch, tuple) or isinstance(data_batch, list):\n",
    "                X_batch, y_batch = data_batch\n",
    "                if data_loader_apply is not None:\n",
    "                    X_batch, y_batch = data_loader_apply(X_batch, y_batch)\n",
    "            else:\n",
    "                X_batch, y_batch = data_loader_apply(data_batch)\n",
    "            X_all.append(X_batch)\n",
    "            y_all.append(y_batch)\n",
    "            if max_validation_iter is not None and i >= max_validation_iter:\n",
    "                break\n",
    "        if not isinstance(X_all[0], torch.Tensor):\n",
    "            X_all = Zip(*X_all, function = torch.cat)\n",
    "        else:\n",
    "            X_all = torch.cat(X_all, 0)\n",
    "        y_all = torch.cat(y_all)\n",
    "        model.plot(X_all, transform_label(y_all))\n",
    "    else:\n",
    "        assert X is not None and y is not None\n",
    "        model.plot(X, transform_label(y))\n",
    "\n",
    "\n",
    "def prepare_inspection(model, data_loader=None, X=None, y=None, transform_label=None, **kwargs):\n",
    "    inspect_functions = kwargs[\"inspect_functions\"] if \"inspect_functions\" in kwargs else None\n",
    "    max_validation_iter = kwargs[\"max_validation_iter\"] if \"max_validation_iter\" in kwargs else None\n",
    "    verbose = kwargs[\"verbose\"] if \"verbose\" in kwargs else False\n",
    "    if transform_label is None:\n",
    "        transform_label = Transform_Label()\n",
    "    if \"loader_process\" in kwargs and kwargs[\"loader_process\"] is not None:\n",
    "        data_loader = kwargs[\"loader_process\"](\"test\")\n",
    "    if data_loader is None:\n",
    "        assert X is not None and y is not None\n",
    "        all_dict_summary = model.prepare_inspection(X, transform_label(y), **kwargs)\n",
    "        if inspect_functions is not None:\n",
    "            for inspect_function_key, inspect_function in inspect_functions.items():\n",
    "                all_dict_summary[inspect_function_key] = inspect_function(model, X, y, **kwargs)\n",
    "    else:\n",
    "        assert X is None and y is None\n",
    "        all_dict = {}\n",
    "        for j, data_batch in enumerate(data_loader):\n",
    "            if verbose is True:\n",
    "                print(\"valid step: {}\".format(j))\n",
    "            if isinstance(data_batch, tuple) or isinstance(data_batch, list):\n",
    "                X_batch, y_batch = data_batch\n",
    "                if \"data_loader_apply\" in kwargs and kwargs[\"data_loader_apply\"] is not None:\n",
    "                    X_batch, y_batch = kwargs[\"data_loader_apply\"](X_batch, y_batch)\n",
    "            else:\n",
    "                X_batch, y_batch = kwargs[\"data_loader_apply\"](data_batch)\n",
    "            info_dict = model.prepare_inspection(X_batch, transform_label(y_batch), valid_step=j, **kwargs)\n",
    "            for key, item in info_dict.items():\n",
    "                if key not in all_dict:\n",
    "                    all_dict[key] = [item]\n",
    "                else:\n",
    "                    all_dict[key].append(item)\n",
    "            if inspect_functions is not None:\n",
    "                for inspect_function_key, inspect_function in inspect_functions.items():\n",
    "                    inspect_function_result = inspect_function(model, X_batch, transform_label(y_batch), **kwargs)\n",
    "                    if inspect_function_key not in all_dict:\n",
    "                        all_dict[inspect_function_key] = [inspect_function_result]\n",
    "                    else:\n",
    "                        all_dict[inspect_function_key].append(inspect_function_result)\n",
    "            if max_validation_iter is not None and j >= max_validation_iter:\n",
    "                break\n",
    "        all_dict_summary = {}\n",
    "        for key, item in all_dict.items():\n",
    "            all_dict_summary[key + \"_val\"] = np.mean(all_dict[key])\n",
    "    return all_dict_summary\n",
    "\n",
    "\n",
    "def get_inspect_items_train(inspect_items):\n",
    "    if inspect_items is None:\n",
    "        return []\n",
    "    inspect_items_train = []\n",
    "    for item in inspect_items:\n",
    "        if item.endswith(\"_tr\"):\n",
    "            inspect_items_train.append(\"_\".join(item.split(\"_\")[:-1]))\n",
    "    return inspect_items_train\n",
    "\n",
    "\n",
    "def update_key_train(info_dict_train, inspect_items_train):\n",
    "    info_dict_train_new = {}\n",
    "    for key, item in info_dict_train.items():\n",
    "        if key in inspect_items_train:\n",
    "            info_dict_train_new[key + \"_tr\"] = item\n",
    "    return deepcopy(info_dict_train_new)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simplification functionality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def simplify(\n",
    "    model,\n",
    "    X=None,\n",
    "    y=None,\n",
    "    mode=\"full\",\n",
    "    isplot=False,\n",
    "    target_name=None,\n",
    "    validation_data=None,\n",
    "    **kwargs\n",
    "):\n",
    "    \"\"\"Simplify a neural network model in various ways. \"model\" can be a single model or a ordered list of models\"\"\"\n",
    "    verbose = kwargs[\"verbose\"] if \"verbose\" in kwargs else 1\n",
    "    if validation_data is None:\n",
    "        X_valid, y_valid = X, y\n",
    "    else:\n",
    "        X_valid, y_valid = validation_data\n",
    "    simplify_criteria = kwargs[\"simplify_criteria\"] if \"simplify_criteria\" in kwargs else (\"DLs\", 0.05, 3, \"relative\") # the first argument choose from \"DL\", \"loss\"\n",
    "    simplify_epsilon = simplify_criteria[1]\n",
    "    simplify_patience = simplify_criteria[2]\n",
    "    simplify_compare_mode = simplify_criteria[3]\n",
    "    performance_monitor = Performance_Monitor(patience = simplify_patience, epsilon = simplify_epsilon, compare_mode = simplify_compare_mode)\n",
    "    record_keys = kwargs[\"record_keys\"] if \"record_keys\" in kwargs else [\"mse\"]\n",
    "    loss_precision_floor = kwargs[\"loss_precision_floor\"] if \"loss_precision_floor\" in kwargs else PrecisionFloorLoss\n",
    "    if X is not None:\n",
    "        if y is None:\n",
    "            y = Variable(forward(model, X, **kwargs).data, requires_grad = False)\n",
    "    if not (isinstance(model, list) or isinstance(model, tuple)):\n",
    "        model = [model]\n",
    "        is_list = False\n",
    "    else:\n",
    "        is_list = True\n",
    "    if mode == \"full\":\n",
    "        mode = [\"collapse_layers\", \"snap\"]\n",
    "    if not isinstance(mode, list):\n",
    "        mode = [mode]\n",
    "\n",
    "    # Obtain the original loss and setup criterion:\n",
    "    loss_type = kwargs[\"loss_type\"] if \"loss_type\" in kwargs else \"mse\"\n",
    "    criterion = get_criterion(loss_type, loss_precision_floor = loss_precision_floor)\n",
    "    DL_criterion = Loss_Fun(core = \"DLs\", loss_precision_floor = loss_precision_floor, DL_sum = True)\n",
    "    loss_dict = OrderedDict()\n",
    "\n",
    "    for mode_ele in mode:\n",
    "        if verbose >= 1:\n",
    "            print(\"\\n\" + \"=\" * 48 + \"\\nSimplifying mode: {}\".format(mode_ele), end = \"\")\n",
    "            if mode_ele == \"snap\":\n",
    "                snap_mode = kwargs[\"snap_mode\"] if \"snap_mode\" in kwargs else \"integer\"\n",
    "                print(\" {}\".format(snap_mode), end = \"\")\n",
    "            if target_name is not None:\n",
    "                print(\" for {}\".format(target_name))\n",
    "            else:\n",
    "                print()\n",
    "            print(\"=\" * 48)\n",
    "        \n",
    "        # Record the loss before simplification:\n",
    "        if X is not None:\n",
    "            pred_valid = forward(model, X_valid, **kwargs)\n",
    "            loss_original = to_np_array(criterion(pred_valid, y_valid))\n",
    "            loss_list = [loss_original]\n",
    "            if verbose >= 1:\n",
    "                print(\"original_loss: {}\".format(loss_original))\n",
    "            mse_record_whole = [to_np_array(nn.MSELoss()(pred_valid, y_valid))]\n",
    "            data_DL_whole = [to_np_array(DL_criterion(pred_valid, y_valid))]\n",
    "        model_DL_whole = [get_model_DL(model)]\n",
    "        event_list = [\"before simplification\"]\n",
    "        iter_end_whole = [1]\n",
    "        is_accept_whole = []\n",
    "        if \"param\" in record_keys:\n",
    "            param_record_whole = [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\")]\n",
    "        if \"param_grad\" in record_keys:\n",
    "            param_grad_record_whole = [model[0].get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True)]\n",
    "        \n",
    "        # Begin simplification:\n",
    "        if mode_ele == \"collapse_layers\":\n",
    "            all_collapse_dict = {}\n",
    "            for model_id, model_ele in enumerate(model):\n",
    "                # Obtain activations for each layer:\n",
    "                activation_list = []\n",
    "                for k in range(len(model_ele.struct_param)):\n",
    "                    if \"activation\" in model_ele.struct_param[k][2]:\n",
    "                        activation_list.append(model_ele.struct_param[k][2][\"activation\"])\n",
    "                    elif \"activation\" in model_ele.settings:\n",
    "                        activation_list.append(model_ele.settings[\"activation\"])\n",
    "                    else:\n",
    "                        activation_list.append(\"default\")\n",
    "                \n",
    "                # Build the collapse_list that stipulates which layers to collapse:\n",
    "                collapse_dict = {}\n",
    "                current_start = None\n",
    "                current_layer_type = None\n",
    "                for k, activation in enumerate(activation_list):\n",
    "                    if activation == \"linear\" and k != len(activation_list) - 1:\n",
    "                        if k not in collapse_dict and current_start is None:\n",
    "                            # Create a new bunch:\n",
    "                            if model_ele.struct_param[k + 1][1] == model_ele.struct_param[k][1]: # The current layer must have the same layer_type as the next layer\n",
    "                                current_start = k\n",
    "                                collapse_dict[current_start] = [k]\n",
    "                                current_layer_type = model_ele.struct_param[k][1]\n",
    "                        else:\n",
    "                            # Adding to current bunch:\n",
    "                            if model_ele.struct_param[k + 1][1] == model_ele.struct_param[k][1] == current_layer_type:\n",
    "                                collapse_dict[current_start].append(k)\n",
    "                            else:\n",
    "                                collapse_dict[current_start].append(k)\n",
    "                                current_start = None\n",
    "                    else:\n",
    "                        if current_start is not None:\n",
    "                            collapse_dict[current_start].append(k)\n",
    "                        current_start = None\n",
    "\n",
    "                # Build new layer:\n",
    "                new_layer_info = {}\n",
    "                for current_start, layer_ids in collapse_dict.items():\n",
    "                    for i, layer_id in enumerate(layer_ids):\n",
    "                        layer = getattr(model_ele, \"layer_{}\".format(layer_id))\n",
    "                        if i == 0:\n",
    "                            W_accum = layer.W_core\n",
    "                            b_accum = layer.b_core\n",
    "                        else:\n",
    "                            W_accum = torch.matmul(W_accum, layer.W_core)\n",
    "                            b_accum = torch.matmul(b_accum, layer.W_core) + layer.b_core\n",
    "                    if model_ele.is_cuda:\n",
    "                        W_accum = W_accum.cpu()\n",
    "                        b_accum = b_accum.cpu()\n",
    "                    last_layer_id = collapse_dict[current_start][-1]\n",
    "                    new_layer_info[current_start] = {\"W_init\": W_accum.data.numpy(), \"b_init\": b_accum.data.numpy(),\n",
    "                                                     \"layer_struct_param\": [b_accum.size(0), model_ele.struct_param[last_layer_id][1], deepcopy(model_ele.struct_param[last_layer_id][2])],\n",
    "                                                    }\n",
    "                    new_layer_info[current_start].pop(\"snap_dict\", None)\n",
    "                if verbose >= 1:\n",
    "                    print(\"model_id {}, layers collapsed: {}\".format(model_id, collapse_dict))\n",
    "                \n",
    "                # Rebuild the Net:\n",
    "                if len(collapse_dict) > 0:\n",
    "                    all_collapse_dict[model_id] = {\"collapse_dict\": collapse_dict, \n",
    "                                                   \"new_layer_info\": new_layer_info, \n",
    "                                                   \"collapse_layer_ids\": [idx for item in collapse_dict.values() for idx in item],\n",
    "                                                  }\n",
    "\n",
    "            # Rebuild the list of models:\n",
    "            if len(all_collapse_dict) > 0:\n",
    "                model_new = []\n",
    "                for model_id, model_ele in enumerate(model):\n",
    "                    if model_id in all_collapse_dict:\n",
    "                        W_list, b_list = model_ele.get_weights_bias(W_source = \"core\", b_source = \"core\")\n",
    "                        W_init_list = []\n",
    "                        b_init_list = []\n",
    "                        struct_param = []\n",
    "                        for k in range(len(model_ele.struct_param)):\n",
    "                            if k not in all_collapse_dict[model_id][\"collapse_layer_ids\"]:\n",
    "                                struct_param.append(model_ele.struct_param[k])\n",
    "                                W_init_list.append(W_list[k])\n",
    "                                b_init_list.append(b_list[k])\n",
    "                            else:\n",
    "                                if k in all_collapse_dict[model_id][\"collapse_dict\"].keys():\n",
    "                                    struct_param.append(all_collapse_dict[model_id][\"new_layer_info\"][k][\"layer_struct_param\"])\n",
    "                                    W_init_list.append(all_collapse_dict[model_id][\"new_layer_info\"][k][\"W_init\"])\n",
    "                                    b_init_list.append(all_collapse_dict[model_id][\"new_layer_info\"][k][\"b_init\"])\n",
    "                        model_ele_new = MLP(input_size = model_ele.input_size,\n",
    "                                            struct_param = struct_param,\n",
    "                                            W_init_list = W_init_list,\n",
    "                                            b_init_list = b_init_list,\n",
    "                                            settings = model_ele.settings,\n",
    "                                            is_cuda = model_ele.is_cuda,\n",
    "                                           )\n",
    "                    else:\n",
    "                        model_ele_new = model_ele\n",
    "                    model_new.append(model_ele_new)               \n",
    "                model = model_new\n",
    "\n",
    "                # Calculate the loss again:\n",
    "                pred_valid = forward(model, X_valid, **kwargs)\n",
    "                loss_new = to_np_array(criterion(pred_valid, y_valid))\n",
    "                if verbose >= 1:\n",
    "                    print(\"after collapsing linear layers in all models, new loss {}\".format(loss_new))\n",
    "                loss_list.append(loss_new)\n",
    "                mse_record_whole.append(to_np_array(nn.MSELoss()(pred_valid, y_valid)))\n",
    "                data_DL_whole.append(to_np_array(DL_criterion(pred_valid, y_valid)))\n",
    "                model_DL_whole.append(get_model_DL(model))\n",
    "                if \"param\" in record_keys:\n",
    "                    param_record_whole.append(model[0].get_weights_bias(W_source = \"core\", b_source = \"core\"))\n",
    "                if \"param_grad\" in record_keys:\n",
    "                    param_grad_record_whole.append(model[0].get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True))\n",
    "                iter_end_whole.append(1)\n",
    "                event_list.append({mode_ele: all_collapse_dict})\n",
    "\n",
    "        elif mode_ele in [\"local\", \"snap\"]:\n",
    "            # 'local': greedily try reducing the input dimension by removing input dimension from the beginning;\n",
    "            # 'snap': greedily snap each float parameter into an integer or rational number. Set argument 'snap_mode' == 'integer' or 'rational'.\n",
    "            if mode_ele == \"snap\":\n",
    "                target_params = [[(model_id, layer_id), \"snap\"] for model_id, model_ele in enumerate(model) for layer_id in range(len(model_ele.struct_param))]\n",
    "            elif mode_ele == \"local\":\n",
    "                for model_id, model_ele in enumerate(model):\n",
    "                    if len(model_ele.struct_param) > 0:\n",
    "                        first_model_id = model_id\n",
    "                        break\n",
    "                first_layer = getattr(model[first_model_id], \"layer_0\")\n",
    "                target_params = [[(first_model_id, 0), [[((\"weight\", (i, j)), 0.) for j in range(first_layer.output_size)]                                                             for i in range(first_layer.input_size)]]]\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "            excluded_idx_dict = {item[0]: [] for item in target_params}\n",
    "            target_layer_ids_exclude = []\n",
    "            for (model_id, layer_id), target_list in target_params:\n",
    "                layer = getattr(model[model_id], \"layer_{}\".format(layer_id))\n",
    "                if isinstance(target_list, list):\n",
    "                    max_passes = len(target_list)\n",
    "                elif target_list == \"snap\":\n",
    "                    max_passes = (layer.input_size + 1) * layer.output_size\n",
    "                    if \"max_passes\" in kwargs:\n",
    "                        max_passes = min(max_passes, kwargs[\"max_passes\"])\n",
    "                else:\n",
    "                    raise Exception(\"target_list {} not recognizable!\".format(target_list))\n",
    "                if verbose >= 2:\n",
    "                    print(\"\\n****starting model:****\")\n",
    "                    model[model_id].get_weights_bias(W_source = \"core\", b_source = \"core\", verbose = True)\n",
    "                    print(\"********\\n\" )\n",
    "                \n",
    "                \n",
    "                performance_monitor.reset()\n",
    "                criteria_value, criteria_result = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                to_stop, pivot_dict, log, _, pivot_id = performance_monitor.monitor(criteria_value, model_dict = model[model_id].model_dict, criteria_result = criteria_result)\n",
    "                for i in range(max_passes):\n",
    "                    # Perform tentative simplification\n",
    "                    if isinstance(target_list, list):\n",
    "                        info = layer.simplify(mode = \"snap\", excluded_idx = excluded_idx_dict[(model_id, layer_id)], snap_targets = target_list[i], **kwargs)\n",
    "                    else:\n",
    "                        info = layer.simplify(mode = \"snap\", excluded_idx = excluded_idx_dict[(model_id, layer_id)], **kwargs)\n",
    "                    if len(info) == 0:\n",
    "                        target_layer_ids_exclude.append((model_id, layer_id))\n",
    "                        print(\"Pass {0}, (model {1}, layer {2}) has no parameters to snap. Revert to pivot model. Go to next layer\".format(i, model_id, layer_id))\n",
    "                        break\n",
    "                    excluded_idx_dict[(model_id, layer_id)] = excluded_idx_dict[(model_id, layer_id)] + info\n",
    "\n",
    "                    _, loss_new, data_record = train_simple(model, X, y, optim_type = \"adam\", validation_data = validation_data, **kwargs)\n",
    "                    if verbose >= 2:\n",
    "                        print(\"=\" * 8)\n",
    "                        model[model_id].get_weights_bias(W_source = \"core\", b_source = \"core\", verbose = True) \n",
    "                    criteria_value, criteria_result = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                    to_stop, pivot_dict, log, is_accept, pivot_id = performance_monitor.monitor(criteria_value, model_dict = model[model_id].model_dict, criteria_result = criteria_result)\n",
    "                    is_accept_whole.append(is_accept)\n",
    "                    if is_accept:\n",
    "                        print('[Accepted] as pivot model!')\n",
    "                        print()\n",
    "\n",
    "                    # Check if the criterion after simplification and refit is worse. If it is worse than the simplify_epsilon, revert:\n",
    "                    if to_stop:\n",
    "                        target_layer_ids_exclude.append((model_id, layer_id))\n",
    "                        if verbose >= 1:\n",
    "                            print(\"Pass {0}, loss: {1}\\tDL: {2}. New snap {3} is do not improve by {4} = {5} for {6} steps. Revert the simplification to pivot model. Go to next layer.\".format(\n",
    "                                i, view_item(log, (\"criteria_result\", \"loss\")), view_item(log, (\"criteria_result\", \"DL\")), info, simplify_criteria[0], simplify_epsilon, simplify_patience))\n",
    "                        break\n",
    "                    mse_record_whole += data_record[\"mse\"]\n",
    "                    data_DL_whole += data_record[\"data_DL\"]\n",
    "                    model_DL_whole += data_record[\"model_DL\"]\n",
    "                    if \"param\" in record_keys:\n",
    "                        param_record_whole += data_record[\"param\"]\n",
    "                    if \"param_grad\" in record_keys:\n",
    "                        param_grad_record_whole += data_record[\"param_grad\"]\n",
    "                    iter_end_whole.append(len(data_record[\"mse\"]))\n",
    "\n",
    "                    model[model_id].reset_layer(layer_id, layer)\n",
    "                    loss_list.append(loss_new)\n",
    "                    event_list.append({mode_ele: ((model_id, layer_id), info)})\n",
    "                    if verbose >= 1:\n",
    "                        print(\"Pass {0}, snap (model {1}, layer {2}), snap {3}. \\tloss: {4}\\tDL: {5}\".format(\n",
    "                            i, model_id, layer_id, info, view_item(log, (\"criteria_result\", \"loss\")), view_item(log, (\"criteria_result\", \"DL\"))))\n",
    "\n",
    "                # Update the whole model's struct_param and snap_dict:\n",
    "                model[model_id].load_model_dict(pivot_dict[\"model_dict\"])\n",
    "                model[model_id].synchronize_settings()\n",
    "                if verbose >= 2:\n",
    "                    print(\"\\n****pivot model at {}th transformation:****\".format(pivot_id))\n",
    "                    model[model_id].get_weights_bias(W_source = \"core\", b_source = \"core\", verbose = True)\n",
    "                    print(\"********\\n\" )\n",
    "\n",
    "        elif mode_ele == \"pair_snap\":\n",
    "            model_new = []\n",
    "            for model_id, model_ele in enumerate(model):\n",
    "                for layer_id, layer_struct_param in enumerate(model_ele.struct_param):\n",
    "                    if layer_struct_param[1] == \"Symbolic_Layer\":\n",
    "                        layer = getattr(model_ele, \"layer_{}\".format(layer_id))\n",
    "                        max_passes = len(layer.get_param_dict()) - 1\n",
    "                        if \"max_passes\" in kwargs:\n",
    "                            max_passes = min(max_passes, kwargs[\"max_passes\"])\n",
    "                        if verbose > 1:\n",
    "                            print(\"original:\")\n",
    "                            print(\"symbolic_expression: \", layer.symbolic_expression)\n",
    "                            print(\"numerical_expression: \", layer.numerical_expression)\n",
    "                            print()\n",
    "\n",
    "                        performance_monitor.reset()\n",
    "                        criteria_value, criteria_result = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                        to_stop, pivot_dict, log, _, pivot_id = performance_monitor.monitor(criteria_value, model_dict = model[model_id].model_dict, criteria_result = criteria_result)\n",
    "                        for i in range(max_passes):\n",
    "                            info = layer.simplify(mode = \"pair_snap\", **kwargs)\n",
    "                            if len(info) == 0:\n",
    "                                target_layer_ids_exclude.append((model_id, layer_id))\n",
    "                                print(\"Pass {0}, (model {1}, layer {2}) has no parameters to pair_snap. Revert to pivot model. Go to next layer\".format(i, model_id, layer_id))\n",
    "                                break\n",
    "                            _, loss, data_record = train_simple(model, X, y, optim_type = \"adam\", epochs = 1000, validation_data = validation_data, **kwargs)\n",
    "                            criteria_value, criteria_result = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                            to_stop, pivot_dict, log, is_accept, pivot_id = performance_monitor.monitor(criteria_value, model_dict = model[model_id].model_dict, criteria_result = criteria_result)\n",
    "                            is_accept_whole.append(is_accept)\n",
    "                            if to_stop:\n",
    "                                if verbose >= 1:\n",
    "                                    print(\"\\nPass {0}, loss: {1}\\tDL: {2}. New snap {3} is do not improve by {4} = {5} for {6} steps. Revert the simplification to pivot model. Go to next layer.\".format(\n",
    "                                        i, view_item(log, (\"criteria_result\", \"loss\")), view_item(log, (\"criteria_result\", \"DL\")), info, simplify_criteria[0], simplify_epsilon, simplify_patience))\n",
    "                                break\n",
    "\n",
    "                            mse_record_whole += data_record[\"mse\"]\n",
    "                            data_DL_whole += data_record[\"data_DL\"]\n",
    "                            model_DL_whole += data_record[\"model_DL\"]\n",
    "                            if \"param\" in record_keys:\n",
    "                                param_record_whole += data_record[\"param\"]\n",
    "                            if \"param_grad\" in record_keys:\n",
    "                                param_grad_record_whole += data_record[\"param_grad\"]\n",
    "                            iter_end_whole.append(len(data_record[\"mse\"]))\n",
    "\n",
    "                            model[model_id].reset_layer(layer_id, layer)\n",
    "                            loss_list.append(loss)\n",
    "                            event_list.append({mode_ele: ((model_id, layer_id), info)})\n",
    "                            if verbose >= 1:\n",
    "                                print(\"\\nPass {0}, snap (model {1}, layer {2}), snap {3}. \\tloss: {4}\\tDL: {5}\".format(\n",
    "                                    i, model_id, layer_id, info, view_item(log, (\"criteria_result\", \"loss\")), view_item(log, (\"criteria_result\", \"DL\"))))\n",
    "                                print(\"symbolic_expression: \", layer.symbolic_expression)\n",
    "                                print(\"numerical_expression: \", layer.numerical_expression)\n",
    "                                print()\n",
    "\n",
    "                        model[model_id].load_model_dict(pivot_dict[\"model_dict\"])\n",
    "                        print(\"final: \\nsymbolic_expression: \", getattr(model[model_id], \"layer_{0}\".format(layer_id)).symbolic_expression)\n",
    "                        print(\"numerical_expression: \", getattr(model[model_id], \"layer_{0}\".format(layer_id)).numerical_expression)\n",
    "                        print()\n",
    "\n",
    "        elif mode_ele[:11] == \"to_symbolic\":\n",
    "            from sympy import Symbol\n",
    "            force_simplification = kwargs[\"force_simplification\"] if \"force_simplification\" in kwargs else False\n",
    "            is_multi_model = True if len(model) > 1 else False\n",
    "            for model_id, model_ele in enumerate(model):\n",
    "                for layer_id, layer_struct_param in enumerate(model_ele.struct_param):\n",
    "                    prefix = \"L{}_\".format(layer_id)\n",
    "                    if layer_struct_param[1] == \"Simple_Layer\":\n",
    "                        # Obtain loss before simplification:\n",
    "                        layer = getattr(model_ele, \"layer_{}\".format(layer_id))\n",
    "                        if X is not None:\n",
    "                            criteria_prev, criteria_result_prev = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                        \n",
    "                        if mode_ele.split(\"_\")[-1] == \"separable\":\n",
    "                            new_layer = Simple_2_Symbolic(layer, settings = model_ele.settings, mode = \"separable\", prefix = prefix)\n",
    "                        else:\n",
    "                            new_layer = Simple_2_Symbolic(layer, settings = model_ele.settings, prefix = prefix)\n",
    "                        model[model_id].reset_layer(layer_id, new_layer)\n",
    "\n",
    "                        if \"snap_dict\" in model_ele.settings and layer_id in model_ele.settings[\"snap_dict\"]:\n",
    "                            subs_targets = []\n",
    "                            for (pos, true_idx), item in model_ele.settings[\"snap_dict\"][layer_id].items():\n",
    "                                if pos == \"weight\":\n",
    "                                    subs_targets.append((Symbol(\"W{0}{1}\".format(true_idx[0], true_idx[1])), item[\"new_value\"]))\n",
    "                                elif pos == \"bias\":\n",
    "                                    subs_targets.append((Symbol(\"b{}\".format(true_idx)), item[\"new_value\"]))\n",
    "                                else:\n",
    "                                    raise Exception(\"pos {} not recognized!\".format(pos))\n",
    "                            new_expression = [expression.subs(subs_targets) for expression in new_layer.symbolic_expression]\n",
    "                            new_layer.set_symbolic_expression(new_expression)\n",
    "                            model_ele.settings[\"snap_dict\"].pop(layer_id)\n",
    "                            model_ele.struct_param[layer_id][2].update(new_layer.struct_param[2])\n",
    "                        \n",
    "                        # Calculate the loss again:\n",
    "                        if X is not None:\n",
    "                            criteria_new, criteria_result_new = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                            if verbose >= 1:\n",
    "                                print(\"Prev_loss: {0}, new loss: {1}\\tprev_DL: {2:.9f}, new DL: {3:.9f}\".format(\n",
    "                                       criteria_result_prev[\"loss\"], criteria_result_new[\"loss\"], criteria_result_prev[\"DL\"], criteria_result_new[\"DL\"]))\n",
    "                                print()\n",
    "                            if criteria_new > criteria_prev * (1 + 0.05):\n",
    "                                print(\"to_symbolic DL increase more than 5%! \", end = \"\")\n",
    "                                if not force_simplification:\n",
    "                                    print(\"Reset layer.\")\n",
    "                                    model[model_id].reset_layer(layer_id, layer)\n",
    "                                else:\n",
    "                                    print(\"Nevertheless, force simplification.\")\n",
    "\n",
    "                            loss_list.append(criteria_result_new[\"loss\"])\n",
    "                            print(\"{0} succeed. Prev_loss: {1}\\tnew_loss: {2}\\tprev_DL: {3:.9f}, new_DL: {4:.9f}\".format(\n",
    "                                    mode_ele, criteria_result_prev[\"loss\"], criteria_result_new[\"loss\"],\n",
    "                                    criteria_result_prev[\"DL\"], criteria_result_new[\"DL\"]))\n",
    "                        else:\n",
    "                            print(\"{0} succeed.\".format(mode_ele))\n",
    "                        event_list.append({mode_ele: (model_id, layer_id)})\n",
    "                        \n",
    "                    \n",
    "                    elif layer_struct_param[1] == \"Sneuron_Layer\":\n",
    "                        # Obtain loss before simplification:\n",
    "                        layer = getattr(model_ele, \"layer_{0}\".format(layer_id))\n",
    "                        criteria_prev, criteria_result_prev = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                        \n",
    "                        new_layer = Sneuron_2_Symbolic(layer, prefix = prefix)\n",
    "                        model[model_id].reset_layer(layer_id, new_layer)\n",
    "                        \n",
    "                        # Calculate the loss again:\n",
    "                        criteria_new, criteria_result_new = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                        if verbose >= 1:\n",
    "                            print(\"Prev_loss: {0}, new loss: {1}\\tprev_DL: {2:.9f}, new DL: {3:.9f}\".format(\n",
    "                                   criteria_result_prev[\"loss\"], criteria_result_new[\"loss\"], criteria_result_prev[\"DL\"], criteria_result_new[\"DL\"]))\n",
    "                            print()\n",
    "                        if criteria_new > criteria_prev * (1 + 0.05):  \n",
    "                            print(\"to_symbolic DL increase more than 5%! \", end = \"\")\n",
    "                            if not force_simplification:\n",
    "                                print(\"Reset layer.\")\n",
    "                                model[model_id].reset_layer(layer_id, layer)\n",
    "                            else:\n",
    "                                print(\"Nevertheless, force simplification.\")\n",
    "                        \n",
    "                        loss_list.append(criteria_result_new[\"loss\"])\n",
    "                        event_list.append({mode_ele: (model_id, layer_id)})\n",
    "                        print(\"{0} succeed. Prev_loss: {1}\\tnew_loss: {2}\\tprev_DL: {3:.9f}, new_DL: {4:.9f}\".format(\n",
    "                                mode_ele, criteria_result_prev[\"loss\"], criteria_result_new[\"loss\"],\n",
    "                                criteria_result_prev[\"DL\"], criteria_result_new[\"DL\"]))\n",
    "            if X is not None:\n",
    "                mse_record_whole.append(to_np_array(nn.MSELoss()(pred_valid, y_valid)))\n",
    "                data_DL_whole.append(to_np_array(DL_criterion(pred_valid, y_valid)))\n",
    "            model_DL_whole.append(get_model_DL(model))\n",
    "            if \"param\" in record_keys:\n",
    "                param_record_whole.append(model[0].get_weights_bias(W_source = \"core\", b_source = \"core\"))\n",
    "            if \"param_grad\" in record_keys:\n",
    "                param_grad_record_whole.append(model[0].get_weights_bias(W_source = \"core\", b_source = \"core\", is_grad = True))\n",
    "            iter_end_whole.append(1)\n",
    "\n",
    "        elif mode_ele == \"symbolic_simplification\":\n",
    "            \"\"\"Collapse multi-layer symbolic expression\"\"\"\n",
    "            from sympy import Symbol, Poly, expand, prod\n",
    "            force_simplification = kwargs[\"force_simplification\"] if \"force_simplification\" in kwargs else False\n",
    "            numerical_threshold = kwargs[\"numerical_threshold\"] if \"numerical_threshold\" in kwargs else None\n",
    "            is_numerical = kwargs[\"is_numerical\"] if \"is_numerical\" in kwargs else False\n",
    "            max_poly_degree = kwargs[\"max_poly_degree\"] if \"max_poly_degree\" in kwargs else None\n",
    "            show_before_truncate = kwargs[\"show_before_truncate\"] if \"show_before_truncate\" in kwargs else False\n",
    "            for model_id, model_ele in enumerate(model):\n",
    "                is_all_symbolic = True\n",
    "                for layer_id, layer_struct_param in enumerate(model_ele.struct_param):\n",
    "                    if layer_struct_param[1] != \"Symbolic_Layer\":\n",
    "                        is_all_symbolic = False\n",
    "                if is_all_symbolic:\n",
    "                    criteria_prev, criteria_result_prev = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                    variables = OrderedDict()\n",
    "                    for i in range(model[0].layer_0.input_size):\n",
    "                        variables[\"x{0}\".format(i)] = Symbol(\"x{0}\".format(i))\n",
    "                    expression = list(variables.values())\n",
    "                    param_dict_all = {}\n",
    "\n",
    "                    # Collapse multiple layers:\n",
    "                    for layer_id, layer_struct_param in enumerate(model_ele.struct_param):\n",
    "                        layer = getattr(model_ele, \"layer_{0}\".format(layer_id))\n",
    "                        layer_expression = deepcopy(layer.numerical_expression)\n",
    "                        layer_expression_new = []\n",
    "                        for expr in layer_expression:\n",
    "                            new_expr = expr.subs({\"x{0}\".format(i): \"t{0}\".format(i) for i in range(len(expression))})  # Use a temporary variable to prevent overriding\n",
    "                            new_expr = new_expr.subs({\"t{0}\".format(i): expression[i] for i in range(len(expression))})\n",
    "                            layer_expression_new.append(expand(new_expr))\n",
    "                        expression = layer_expression_new\n",
    "                    \n",
    "                    # Show full expression before performing truncation:\n",
    "                    if show_before_truncate:\n",
    "                        for i, expr in enumerate(expression):\n",
    "                            print(\"Full expression {0}:\".format(i))\n",
    "                            pp.pprint(Poly(expr, *list(variables.values())))\n",
    "                            print()\n",
    "\n",
    "                    model_ele_candidate = MLP(input_size = model[0].layer_0.input_size,\n",
    "                                              struct_param = [[layer.output_size, \"Symbolic_Layer\", {\"symbolic_expression\": \"x0\"}]],\n",
    "                                              settings = {},\n",
    "                                              is_cuda = model_ele.is_cuda,\n",
    "                                             )\n",
    "                    # Setting maximul degree for polynomial:\n",
    "                    if max_poly_degree is not None:\n",
    "                        new_expression = []\n",
    "                        for expr in expression:\n",
    "                            expr = Poly(expr, *list(variables.values()))\n",
    "                            degree_list = []\n",
    "                            coeff_list = []\n",
    "                            for degree, coeff in expr.terms():\n",
    "                                # Only use monomials with degree not larger than max_poly_degree:\n",
    "                                if sum(degree) <= max_poly_degree: \n",
    "                                    degree_list.append(degree)\n",
    "                                    coeff_list.append(coeff)\n",
    "\n",
    "                            new_expr = 0\n",
    "                            for degree, coeff in zip(degree_list, coeff_list):\n",
    "                                new_expr += prod([variables[\"x{0}\".format(i)] ** degree[i] for i in range(len(degree))]) * coeff\n",
    "                            new_expression.append(new_expr)\n",
    "                        expression = new_expression\n",
    "\n",
    "                    # Update symbolic expression for model_ele_candidate:\n",
    "                    if not is_numerical:\n",
    "                        param_dict_all = {}\n",
    "                        expression_new_all = []\n",
    "                        for expr in expression:\n",
    "                            expression_new, param_dict = numerical_2_parameter(expr, idx = len(param_dict_all), threshold = numerical_threshold)\n",
    "                            expression_new_all.append(expression_new)\n",
    "                            param_dict_all.update(param_dict)\n",
    "                        model_ele_candidate.layer_0.set_symbolic_expression(expression_new_all, p_init = param_dict_all)\n",
    "                    else:\n",
    "                        model_ele_candidate.layer_0.set_symbolic_expression(expression)\n",
    "                        model_ele_candidate.layer_0.set_numerical(True)\n",
    "                    \n",
    "                    criteria_new, criteria_result_new = get_criteria_value(model_ele_candidate, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                    if criteria_new > criteria_prev * (1 + 0.05):                            \n",
    "                        print(\"to_symbolic DL increase more than 5%! \", end = \"\")\n",
    "                        if force_simplification:\n",
    "                            print(\"Nevertheless, force simplification.\")\n",
    "                            model[model_id] = model_ele_candidate\n",
    "                        else:\n",
    "                            print(\"Revert.\")\n",
    "                    else:\n",
    "                        model[model_id] = model_ele_candidate\n",
    "\n",
    "        elif mode_ele == \"activation_snap\":\n",
    "            from sympy import Function\n",
    "            def get_sign_snap_candidate(layer, activation_source, excluded_neurons = None):\n",
    "                coeff_dict = {}\n",
    "                for i in range(len(layer.symbolic_expression)):\n",
    "                    current_expression = [layer.symbolic_expression[i]]\n",
    "                    func_names = layer.get_function_name_list(current_expression)\n",
    "                    if activation_source in func_names:\n",
    "                        coeff = [element for element in layer.get_param_name_list(current_expression) if element[0] == \"W\"]\n",
    "                        coeff_dict[i] = np.mean([np.abs(value) for key, value in layer.get_param_dict().items() if key in coeff])\n",
    "                best_idx = None\n",
    "                best_value = 0\n",
    "                for key, value in coeff_dict.items():\n",
    "                    if value > best_value and key not in excluded_neurons:\n",
    "                        best_value = value\n",
    "                        best_idx = key\n",
    "                return best_idx, best_value\n",
    "\n",
    "            activation_source = kwargs[\"activation_source\"] if \"activation_source\" in kwargs else \"sigmoid\"\n",
    "            activation_target = kwargs[\"activation_target\"] if \"activation_target\" in kwargs else \"heaviside\"\n",
    "            activation_fun_source = Function(activation_source)\n",
    "            activation_fun_target = Function(activation_target)\n",
    "\n",
    "            for model_id, model_ele in enumerate(model):\n",
    "                for layer_id, layer_struct_param in enumerate(model_ele.struct_param):\n",
    "                    if layer_struct_param[1] == \"Symbolic_Layer\":\n",
    "                        layer = getattr(model_ele, \"layer_{0}\".format(layer_id))\n",
    "                        excluded_neurons = []\n",
    "                        if activation_source not in layer.get_function_name_list():\n",
    "                            continue\n",
    "\n",
    "                        performance_monitor.reset()\n",
    "                        criteria_value, criteria_result = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                        to_stop, pivot_dict, log, _, pivot_id = performance_monitor.monitor(criteria_value, model_dict = model[model_id].model_dict, criteria_result = criteria_result)\n",
    "                        for i in range(layer_struct_param[0]):\n",
    "                            # Obtain loss before simplification:\n",
    "                            layer = getattr(model_ele, \"layer_{0}\".format(layer_id))\n",
    "                            best_idx, _ = get_sign_snap_candidate(layer, activation_source, excluded_neurons = excluded_neurons)\n",
    "                            excluded_neurons.append(best_idx)\n",
    "                            \n",
    "                            new_expression = [expression.subs(activation_fun_source, activation_fun_target) if j == best_idx else expression for j, expression in enumerate(layer.symbolic_expression)]\n",
    "                            print(\"Pass {0}, candidate new expression: {1}\".format(i, new_expression))\n",
    "                            layer.set_symbolic_expression(new_expression)\n",
    "\n",
    "                            # Train:\n",
    "                            _, loss_new, data_record = train_simple(model, X, y, validation_data = validation_data, **kwargs)\n",
    "\n",
    "                            criteria_value, criteria_result = get_criteria_value(model, X, y, criteria_type = simplify_criteria[0], criterion = criterion, **kwargs)\n",
    "                            to_stop, pivot_dict, log, is_accept, pivot_id = performance_monitor.monitor(criteria_value, model_dict = model[model_id].model_dict, criteria_result = criteria_result)\n",
    "                            is_accept_whole.append(is_accept)\n",
    "                            # Check if the criterion after simplification and refit is worse. If it is worse than the simplify_epsilon, revert:\n",
    "                            if to_stop:\n",
    "                                model[model_id].load_model_dict(pivot_dict[\"model_dict\"])\n",
    "                                if verbose >= 1:\n",
    "                                    print(\"Pass {0}, loss: {1}\\tDL: {2}. New snap {3} is do not improve by {4} = {5} for {6} steps. Revert the simplification to pivot model. Continue\".format(\n",
    "                                        i, view_item(log, (\"criteria_result\", \"loss\")), view_item(log, (\"criteria_result\", \"DL\")), info, simplify_criteria[0], simplify_epsilon, simplify_patience))\n",
    "                                continue   \n",
    "                                \n",
    "                            mse_record_whole += data_record[\"mse\"]\n",
    "                            data_DL_whole += data_record[\"data_DL\"]\n",
    "                            model_DL_whole += data_record[\"model_DL\"]\n",
    "                            if \"param\" in record_keys:\n",
    "                                param_record_whole += data_record[\"param\"]\n",
    "                            if \"param_grad\" in record_keys:\n",
    "                                param_grad_record_whole += data_record[\"param_grad\"]\n",
    "                            iter_end_whole.append(len(data_record[\"mse\"]))\n",
    "\n",
    "                            loss_list.append(loss_new)\n",
    "                            event_list.append({mode_ele: (model_id, layer_id)})\n",
    "                            if verbose >= 1:\n",
    "                                print(\"{0} succeed at (model {1}, layer {2}). loss: {3}\\tDL: {4}\".format(\n",
    "                                    mode_ele, model_id, layer_id, view_item(log, (\"criteria_result\", \"loss\")), view_item(log, (\"criteria_result\", \"DL\"))))\n",
    "                                print(\"symbolic_expression: \", layer.symbolic_expression)\n",
    "                                print(\"numerical_expression: \", layer.numerical_expression)\n",
    "                                print()\n",
    "                        model[model_id].load_model_dict(pivot_dict[\"model_dict\"])\n",
    "        \n",
    "        elif mode_ele == \"ramping-L1\":\n",
    "            loss_list_specific = []\n",
    "            ramping_L1_list = kwargs[\"ramping_L1_list\"] if \"ramping_L1_list\" in kwargs else np.logspace(-7, -1, 30)\n",
    "            ramping_mse_threshold = kwargs[\"ramping_mse_threshold\"] if \"ramping_mse_threshold\" in kwargs else 1e-5\n",
    "            ramping_final_multiplier = kwargs[\"ramping_final_multiplier\"] if \"ramping_final_multiplier\" in kwargs else 1e-2\n",
    "            layer_dict_dict = {}\n",
    "            for i, L1_amp in enumerate(ramping_L1_list):\n",
    "                reg_dict = {\"weight\": L1_amp, \"bias\": L1_amp, \"param\": L1_amp}\n",
    "                _, loss_end, data_record = train_simple(model, X, y, reg_dict = reg_dict, patience = None, validation_data = validation_data, **kwargs)\n",
    "                layer_dict_dict[i] = model[0].layer_0.layer_dict\n",
    "                weight, bias = model[0].layer_0.get_weights_bias()\n",
    "                print(\"L1-amp: {0}\\tloss: {1}\\tweight: {2}\\tbias: {3}\".format(L1_amp, loss_end, weight, bias))\n",
    "                loss_list_specific.append(loss_end)\n",
    "                if \"param\" in record_keys:\n",
    "                    param_record_whole.append((weight, bias))\n",
    "                if loss_end > ramping_mse_threshold:\n",
    "                    if len(loss_list_specific) == 1:\n",
    "                        print(\"\\nThe MSE after the first L1-amp={0} is already larger than the ramping_mse_threshold. Stop and use current L1-amp. The figures will look empty.\".format(ramping_mse_threshold))\n",
    "                    else:\n",
    "                        print(\"\\nThe MSE {0} is larger than the ramping_mse_threshold {1}, stop ramping-L1 simplification\".format(loss_end, ramping_mse_threshold))\n",
    "                    break \n",
    "                mse_record_whole.append(data_record[\"mse\"][-1])\n",
    "                data_DL_whole.append(data_record[\"data_DL\"][-1])\n",
    "                model_DL_whole.append(data_record[\"model_DL\"][-1])\n",
    "                iter_end_whole.append(1)\n",
    "            final_L1_amp = L1_amp * ramping_final_multiplier\n",
    "            final_L1_idx = np.argmin(np.abs(np.array(ramping_L1_list) - final_L1_amp))\n",
    "            layer_dict_final = layer_dict_dict[final_L1_idx]\n",
    "            print(\"Final L1_amp used: {0}\".format(ramping_L1_list[final_L1_idx]))\n",
    "            if \"param\" in record_keys:\n",
    "                print(\"Final param value:\\nweights: {0}\\nbias{1}\".format(param_record_whole[final_L1_idx][0], param_record_whole[final_L1_idx][1]))\n",
    "            model[0].layer_0.load_layer_dict(layer_dict_final)\n",
    "            mse_record_whole = mse_record_whole[: final_L1_idx + 2]\n",
    "            data_DL_whole = data_DL_whole[: final_L1_idx + 2]\n",
    "            model_DL_whole = model_DL_whole[: final_L1_idx + 2]\n",
    "            iter_end_whole = iter_end_whole[: final_L1_idx + 2]\n",
    "\n",
    "            if isplot:\n",
    "                def dict_to_list(Dict):\n",
    "                    return np.array([value for value in Dict.values()])\n",
    "                weights_list = []\n",
    "                bias_list = []\n",
    "                for element in param_record_whole:\n",
    "                    if isinstance(element[0], dict):\n",
    "                        element_core = dict_to_list(element[0])\n",
    "                        weights_list.append(element_core)\n",
    "                    else:\n",
    "                        element_core = to_np_array(element[0]).squeeze(1)\n",
    "                        weights_list.append(element_core)\n",
    "                        bias_list.append(to_np_array(element[1]))\n",
    "                weights_list = np.array(weights_list)\n",
    "                bias_list = np.array(bias_list).squeeze(1)\n",
    "\n",
    "                import matplotlib.pylab as plt\n",
    "                plt.figure(figsize = (7,5))\n",
    "                plt.loglog(ramping_L1_list[: len(loss_list_specific)], loss_list_specific)\n",
    "                plt.xlabel(\"L1 amp\", fontsize = 16)\n",
    "                plt.ylabel(\"mse\", fontsize = 16)\n",
    "                plt.show()\n",
    "\n",
    "                plt.figure(figsize = (7,5))\n",
    "                plt.semilogx(ramping_L1_list[: len(loss_list_specific)], loss_list_specific)\n",
    "                plt.xlabel(\"L1 amp\", fontsize = 16)\n",
    "                plt.ylabel(\"mse\", fontsize = 16)\n",
    "                plt.show()\n",
    "\n",
    "                plt.figure(figsize = (7,5))\n",
    "                for i in range(weights_list.shape[1]):\n",
    "                    plt.semilogx(ramping_L1_list[: len(loss_list_specific)], weights_list[:,i], label = \"weight_{0}\".format(i))\n",
    "                if len(bias_list) > 0:\n",
    "                    plt.semilogx(ramping_L1_list[: len(loss_list_specific)], bias_list, label = \"bias\")\n",
    "                plt.xlabel(\"L1 amp\", fontsize = 16)\n",
    "                plt.ylabel(\"parameter_values\", fontsize = 16)\n",
    "                plt.legend()\n",
    "                plt.show()\n",
    "                plt.clf()\n",
    "                plt.close()\n",
    "        else:\n",
    "            raise Exception(\"mode {0} not recognized!\".format(mode_ele))\n",
    "\n",
    "        loss_dict[mode_ele] = {}\n",
    "        if X is not None:\n",
    "            loss_dict[mode_ele][\"mse_record_whole\"] = mse_record_whole\n",
    "            loss_dict[mode_ele][\"data_DL_whole\"] = data_DL_whole\n",
    "            loss_dict[mode_ele][\"{0}_test\".format(loss_type)] = loss_list\n",
    "        loss_dict[mode_ele][\"model_DL_whole\"] = model_DL_whole\n",
    "        if \"param\" in record_keys:\n",
    "            loss_dict[mode_ele][\"param_record_whole\"] = param_record_whole\n",
    "        if \"param_grad\" in record_keys:\n",
    "            loss_dict[mode_ele][\"param_grad_record_whole\"] = param_grad_record_whole\n",
    "        loss_dict[mode_ele][\"iter_end_whole\"] = iter_end_whole\n",
    "        loss_dict[mode_ele][\"event_list\"] = event_list\n",
    "        loss_dict[mode_ele][\"is_accept_whole\"] = is_accept_whole\n",
    "        if mode_ele == \"ramping-L1\":\n",
    "            loss_dict[mode_ele][\"ramping_L1_list\"] = ramping_L1_list\n",
    "            loss_dict[mode_ele][\"loss_list_specific\"] = loss_list_specific\n",
    "\n",
    "    if not is_list:\n",
    "        model = model[0]\n",
    "    \n",
    "    return model, loss_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model architectures:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### MLP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        struct_param = None,\n",
    "        W_init_list = None,     # initialization for weights\n",
    "        b_init_list = None,     # initialization for bias\n",
    "        settings = {},          # Default settings for each layer, if the settings for the layer is not provided in struct_param\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        super(MLP, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.is_cuda = is_cuda\n",
    "        self.settings = deepcopy(settings)\n",
    "        if struct_param is not None:\n",
    "            self.num_layers = len(struct_param)\n",
    "            self.W_init_list = W_init_list\n",
    "            self.b_init_list = b_init_list\n",
    "            self.info_dict = {}\n",
    "\n",
    "            self.init_layers(deepcopy(struct_param))\n",
    "        else:\n",
    "            self.num_layers = 0\n",
    "\n",
    "\n",
    "    @property\n",
    "    def struct_param(self):\n",
    "        return [getattr(self, \"layer_{0}\".format(i)).struct_param for i in range(self.num_layers)]\n",
    "\n",
    "    \n",
    "    @property\n",
    "    def output_size(self):\n",
    "        return self.get_layer(-1).output_size\n",
    "\n",
    "\n",
    "    @property\n",
    "    def structure(self):\n",
    "        structure = OrderedDict()\n",
    "        structure[\"input_size\"] = self.input_size\n",
    "        structure[\"output_size\"] = self.output_size\n",
    "        structure[\"struct_param\"] = self.struct_param if hasattr(self, \"struct_param\") else None\n",
    "        return structure\n",
    "\n",
    "\n",
    "    def init_layers(self, struct_param):\n",
    "        res_forward = self.settings[\"res_forward\"] if \"res_forward\" in self.settings else False\n",
    "        for k, layer_struct_param in enumerate(struct_param):\n",
    "            if res_forward:\n",
    "                num_neurons_prev = struct_param[k - 1][0] + self.input_size if k > 0 else self.input_size\n",
    "            else:\n",
    "                num_neurons_prev = struct_param[k - 1][0] if k > 0 else self.input_size\n",
    "            num_neurons = layer_struct_param[0]\n",
    "            W_init = self.W_init_list[k] if self.W_init_list is not None else None\n",
    "            b_init = self.b_init_list[k] if self.b_init_list is not None else None\n",
    "\n",
    "            # Get settings for the current layer:\n",
    "            layer_settings = deepcopy(self.settings) if bool(self.settings) else {}\n",
    "            layer_settings.update(layer_struct_param[2])            \n",
    "\n",
    "            # Construct layer:\n",
    "            layer = get_Layer(layer_type = layer_struct_param[1],\n",
    "                              input_size = num_neurons_prev,\n",
    "                              output_size = num_neurons,\n",
    "                              W_init = W_init,\n",
    "                              b_init = b_init,\n",
    "                              settings = layer_settings,\n",
    "                              is_cuda = self.is_cuda,\n",
    "                             )\n",
    "            setattr(self, \"layer_{}\".format(k), layer)\n",
    "\n",
    "\n",
    "    def forward(self, *input, p_dict=None, **kwargs):\n",
    "        kwargs = filter_kwargs(kwargs, [\"res_forward\", \"is_res_block\", \"act_noise_scale\"])  # only allow certain kwargs to be passed\n",
    "        if isinstance(input, tuple):\n",
    "            input = torch.cat(input, -1)\n",
    "        output = input\n",
    "        res_forward = self.settings[\"res_forward\"] if \"res_forward\" in self.settings else False\n",
    "        is_res_block = self.settings[\"is_res_block\"] if \"is_res_block\" in self.settings else False\n",
    "        for k in range(len(self.struct_param)):\n",
    "            p_dict_ele = p_dict[k] if p_dict is not None else None\n",
    "            if res_forward and k > 0:\n",
    "                output = getattr(self, \"layer_{}\".format(k))(torch.cat([output, input], -1), p_dict=p_dict_ele, **kwargs)\n",
    "            else:\n",
    "                output = getattr(self, \"layer_{}\".format(k))(output, p_dict=p_dict_ele, **kwargs)\n",
    "        if is_res_block:\n",
    "            output = output + input\n",
    "        return output\n",
    "    \n",
    "    \n",
    "    def copy(self):\n",
    "        return deepcopy(self)\n",
    "\n",
    "\n",
    "    def simplify(self, X=None, y=None, mode=\"full\", isplot=False, target_name=None, validation_data = None, **kwargs):\n",
    "        new_model, _ = simplify(self, X, y, mode=mode, isplot=isplot, target_name=target_name, validation_data=validation_data, **kwargs)\n",
    "        self.__dict__.update(new_model.__dict__)\n",
    "    \n",
    "    \n",
    "    def snap(self, snap_mode=\"integer\", top=5, **kwargs):\n",
    "        \"\"\"Generate a set of new models whose parameters are snapped, each model with a different number of snapped parameters.\"\"\"\n",
    "        if not hasattr(self, \"num_layers\") or self.num_layers != 1:\n",
    "            return False, [self]\n",
    "        else:\n",
    "            model_list = []\n",
    "            top = top if snap_mode != \"unsnap\" else 1\n",
    "            for top_ele in range(1, top + 1):\n",
    "                new_model = self.copy()\n",
    "                layer = new_model.layer_0\n",
    "                info_list = layer.simplify(mode=\"snap\", top=top_ele, snap_mode=snap_mode)\n",
    "                if len(info_list) > 0:\n",
    "                    new_model.reset_layer(0, layer)\n",
    "                    model_list.append(new_model)\n",
    "            is_succeed = len(model_list) > 0\n",
    "            return is_succeed, model_list\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        reg = to_Variable([0], is_cuda=self.is_cuda)\n",
    "        for k in range(len(self.struct_param)):\n",
    "            layer = getattr(self, \"layer_{}\".format(k))\n",
    "            reg = reg + layer.get_regularization(mode = mode, source = source)\n",
    "        return reg\n",
    "\n",
    "\n",
    "    def get_layer(self, layer_id):\n",
    "        if layer_id < 0:\n",
    "            layer_id += self.num_layers\n",
    "        return getattr(self, \"layer_{}\".format(layer_id))\n",
    "\n",
    "\n",
    "    def reset_layer(self, layer_id, layer):\n",
    "        setattr(self, \"layer_{}\".format(layer_id), layer)\n",
    "\n",
    "\n",
    "    def insert_layer(self, layer_id, layer):\n",
    "        if layer_id < 0:\n",
    "            layer_id += self.num_layers\n",
    "        if layer_id < self.num_layers - 1:\n",
    "            next_layer = getattr(self, \"layer_{}\".format(layer_id + 1))\n",
    "            if next_layer.struct_param[1] == \"Simple_Layer\":\n",
    "                assert next_layer.input_size == layer.output_size, \"The inserted layer's output_size {0} must be compatible with next layer_{1}'s input_size {2}!\"\\\n",
    "                    .format(layer.output_size, layer_id + 1, next_layer.input_size)\n",
    "        for i in range(self.num_layers - 1, layer_id - 1, -1):\n",
    "            setattr(self, \"layer_{}\".format(i + 1), getattr(self, \"layer_{}\".format(i)))\n",
    "        setattr(self, \"layer_{}\".format(layer_id), layer)\n",
    "        self.num_layers += 1\n",
    "    \n",
    "    \n",
    "    def remove_layer(self, layer_id):\n",
    "        if layer_id < 0:\n",
    "            layer_id += self.num_layers\n",
    "        if layer_id < self.num_layers - 1:\n",
    "            num_neurons_prev = self.struct_param[layer_id - 1][0] if layer_id > 0 else self.input_size\n",
    "            replaced_layer = getattr(self, \"layer_{}\".format(layer_id + 1))\n",
    "            if replaced_layer.struct_param[1] == \"Simple_Layer\":\n",
    "                assert replaced_layer.input_size == num_neurons_prev, \\\n",
    "                    \"After deleting layer_{0}, the replaced layer's input_size {1} must be compatible with previous layer's output neurons {2}!\"\\\n",
    "                        .format(layer_id, replaced_layer.input_size, num_neurons_prev)\n",
    "        for i in range(layer_id, self.num_layers - 1):\n",
    "            setattr(self, \"layer_{}\".format(i), getattr(self, \"layer_{}\".format(i + 1)))\n",
    "        self.num_layers -= 1\n",
    "\n",
    "\n",
    "    def prune_neurons(self, layer_id, neuron_ids):\n",
    "        if layer_id == \"input\":\n",
    "            layer = self.get_layer(0)\n",
    "            layer.prune_input_neurons(neuron_ids)\n",
    "            self.input_size = layer.input_size\n",
    "        else:\n",
    "            if layer_id < 0:\n",
    "                layer_id = self.num_layers + layer_id\n",
    "            layer = getattr(self, \"layer_{}\".format(layer_id))\n",
    "            layer.prune_output_neurons(neuron_ids)\n",
    "            self.reset_layer(layer_id, layer)\n",
    "            if layer_id < self.num_layers - 1:\n",
    "                next_layer = getattr(self, \"layer_{}\".format(layer_id + 1))\n",
    "                next_layer.prune_input_neurons(neuron_ids)\n",
    "                self.reset_layer(layer_id + 1, next_layer)\n",
    "\n",
    "\n",
    "    def add_neurons(self, layer_id, num_neurons, mode = (\"imitation\", \"zeros\")):\n",
    "        if not isinstance(mode, list) and not isinstance(mode, tuple):\n",
    "            mode = (mode, mode)\n",
    "        if layer_id < 0:\n",
    "            layer_id = self.num_layers + layer_id\n",
    "        layer = getattr(self, \"layer_{}\".format(layer_id))\n",
    "        layer.add_output_neurons(num_neurons, mode = mode[0])\n",
    "        self.reset_layer(layer_id, layer)\n",
    "        if layer_id < self.num_layers - 1:\n",
    "            next_layer = getattr(self, \"layer_{}\".format(layer_id + 1))\n",
    "            next_layer.add_input_neurons(num_neurons, mode = mode[1])\n",
    "            self.reset_layer(layer_id + 1, next_layer)\n",
    "        if layer_id == 0:\n",
    "            self.input_size = self.get_layer(0).input_size\n",
    "\n",
    "    \n",
    "    def inspect_operation(self, input, operation_between, p_dict = None, **kwargs):\n",
    "        output = input\n",
    "        res_forward = self.settings[\"res_forward\"] if \"res_forward\" in self.settings else False\n",
    "        is_res_block = self.settings[\"is_res_block\"] if \"is_res_block\" in self.settings else False\n",
    "        for k in range(*operation_between):\n",
    "            p_dict_ele = p_dict[k] if p_dict is not None else None\n",
    "            if res_forward and k > 0:\n",
    "                output = getattr(self, \"layer_{}\".format(k))(torch.cat([output, input], -1), p_dict = p_dict_ele)\n",
    "            else:\n",
    "                output = getattr(self, \"layer_{}\".format(k))(output, p_dict = p_dict_ele)\n",
    "        if is_res_block:\n",
    "            output = output + input\n",
    "        return output\n",
    "\n",
    "\n",
    "    def get_weights_bias(self, W_source = \"core\", b_source = \"core\", layer_ids = None, is_grad = False, isplot = False, verbose = False, raise_error = True):\n",
    "        if not hasattr(self, \"struct_param\"):\n",
    "            return None, None\n",
    "        layer_ids = range(len(self.struct_param)) if layer_ids is None else layer_ids\n",
    "        W_list = []\n",
    "        b_list = []\n",
    "        if W_source is not None:\n",
    "            for k in range(len(self.struct_param)):\n",
    "                if k in layer_ids:\n",
    "                    if W_source == \"core\":\n",
    "                        try:\n",
    "                            W, _ = getattr(self, \"layer_{}\".format(k)).get_weights_bias(is_grad = is_grad)\n",
    "                        except Exception as e:\n",
    "                            if raise_error:\n",
    "                                raise\n",
    "                            else:\n",
    "                                print(e)\n",
    "                            W = np.array([np.NaN])\n",
    "                    else:\n",
    "                        raise Exception(\"W_source '{}' not recognized!\".format(W_source))\n",
    "                    W_list.append(W)\n",
    "        \n",
    "        if b_source is not None:\n",
    "            for k in range(len(self.struct_param)):\n",
    "                if k in layer_ids:\n",
    "                    if b_source == \"core\":\n",
    "                        try:\n",
    "                            _, b = getattr(self, \"layer_{}\".format(k)).get_weights_bias(is_grad = is_grad)\n",
    "                        except Exception as e:\n",
    "                            if raise_error:\n",
    "                                raise\n",
    "                            else:\n",
    "                                print(e)\n",
    "                            b = np.array([np.NaN])\n",
    "                    else:\n",
    "                        raise Exception(\"b_source '{}' not recognized!\".format(b_source))\n",
    "                b_list.append(b)\n",
    "\n",
    "        if verbose:\n",
    "            import pprint as pp\n",
    "            if W_source is not None:\n",
    "                print(\"weight:\")\n",
    "                pp.pprint(W_list)\n",
    "            if b_source is not None:\n",
    "                print(\"bias:\")\n",
    "                pp.pprint(b_list)\n",
    "                \n",
    "        if isplot:\n",
    "            if W_source is not None:\n",
    "                print(\"weight {}:\".format(W_source))\n",
    "                plot_matrices(W_list)\n",
    "            if b_source is not None:\n",
    "                print(\"bias {}:\".format(b_source))\n",
    "                plot_matrices(b_list)\n",
    "\n",
    "        return W_list, b_list\n",
    "\n",
    "\n",
    "    def split_to_model_ensemble(self, mode = \"standardize\"):\n",
    "        num_models = self.struct_param[-1][0]\n",
    "        model_core = deepcopy(self)\n",
    "        if mode == \"standardize\":\n",
    "            last_layer = getattr(model_core, \"layer_{}\".format(model_core.num_layers - 1))\n",
    "            last_layer.standardize(mode = \"b_mean_zero\")\n",
    "        else:\n",
    "            raise Exception(\"mode {} not recognized!\".format(mode))\n",
    "        model_list = [deepcopy(model_core) for i in range(num_models)]\n",
    "        for i, model in enumerate(model_list):\n",
    "            to_prune = list(range(num_models))\n",
    "            to_prune.pop(i)\n",
    "            model.prune_neurons(-1, to_prune)\n",
    "        return construct_model_ensemble_from_nets(model_list)\n",
    "    \n",
    "    \n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": self.__class__.__name__}\n",
    "        model_dict[\"input_size\"] = self.input_size\n",
    "        model_dict[\"struct_param\"] = get_full_struct_param(self.struct_param, self.settings)\n",
    "        model_dict[\"weights\"], model_dict[\"bias\"] = self.get_weights_bias(W_source = \"core\", b_source = \"core\")\n",
    "        model_dict[\"settings\"] = deepcopy(self.settings)\n",
    "        model_dict[\"net_type\"] = self.__class__.__name__\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    @property\n",
    "    def DL(self):\n",
    "        return np.sum([getattr(self, \"layer_{}\".format(i)).DL for i in range(self.num_layers)])\n",
    "\n",
    "\n",
    "    def load_model_dict(self, model_dict):\n",
    "        new_net = load_model_dict_net(model_dict, is_cuda = self.is_cuda)\n",
    "        self.__dict__.update(new_net.__dict__)\n",
    "    \n",
    "    \n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "\n",
    "    def get_loss(self, input, target, criterion, **kwargs):\n",
    "        y_pred = self(input, **kwargs)\n",
    "        return criterion(y_pred, target)\n",
    "\n",
    "\n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return {}\n",
    "\n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        for k in range(self.num_layers):\n",
    "            getattr(self, \"layer_{}\".format(k)).set_cuda(is_cuda)\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        for k in range(self.num_layers):\n",
    "            getattr(self, \"layer_{}\".format(k)).set_trainable(is_trainable)\n",
    "\n",
    "\n",
    "    def get_snap_dict(self):\n",
    "        snap_dict = {}\n",
    "        for k in range(len(self.struct_param)):\n",
    "            layer = getattr(self, \"layer_{}\".format(k))\n",
    "            if hasattr(layer, \"snap_dict\"):\n",
    "                recorded_layer_snap_dict = {}\n",
    "                for key, item in layer.snap_dict.items():\n",
    "                    recorded_layer_snap_dict[key] = {\"new_value\": item[\"new_value\"]}\n",
    "                if len(recorded_layer_snap_dict) > 0:\n",
    "                    snap_dict[k] = recorded_layer_snap_dict\n",
    "        return snap_dict\n",
    "\n",
    "\n",
    "    def synchronize_settings(self):\n",
    "        snap_dict = self.get_snap_dict()\n",
    "        if len(snap_dict) > 0:\n",
    "            self.settings[\"snap_dict\"] = snap_dict\n",
    "        return self.settings\n",
    "\n",
    "\n",
    "    def get_sympy_expression(self, verbose = True):\n",
    "        expressions = {i: {} for i in range(self.num_layers)}\n",
    "        for i in range(self.num_layers):\n",
    "            layer = getattr(self, \"layer_{}\".format(i))\n",
    "            if layer.struct_param[1] == \"Symbolic_Layer\":\n",
    "                if verbose:\n",
    "                    print(\"Layer {}, symbolic_expression:  {}\".format(i, layer.symbolic_expression))\n",
    "                    print(\"          numerical_expression: {}\".format(layer.numerical_expression))\n",
    "                expressions[i][\"symbolic_expression\"] = layer.symbolic_expression\n",
    "                expressions[i][\"numerical_expression\"] = layer.numerical_expression\n",
    "                expressions[i][\"param_dict\"] = layer.get_param_dict()\n",
    "                expressions[i][\"DL\"] = layer.DL\n",
    "            else:\n",
    "                if verbose:\n",
    "                    print(\"Layer {} is not a symbolic layer.\".format(i))\n",
    "                expressions[i] = None\n",
    "        return expressions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Labelmix_MLP:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Labelmix_MLP(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        struct_param,\n",
    "        idx_label=None,\n",
    "        is_cuda=False,\n",
    "    ):\n",
    "        super(Labelmix_MLP, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.struct_param = struct_param\n",
    "        self.num_layers = len(struct_param)\n",
    "        if idx_label is not None and len(idx_label) == input_size:\n",
    "            idx_label = None\n",
    "        if idx_label is not None:\n",
    "            self.idx_label = torch.LongTensor(idx_label)\n",
    "            idx_main = list(set(range(input_size)) - set(to_np_array(idx_label).astype(int).tolist()))\n",
    "            self.idx_main = torch.LongTensor(idx_main)\n",
    "        else:\n",
    "            self.idx_label = None\n",
    "            self.idx_main = torch.LongTensor(list(range(input_size)))\n",
    "        num_neurons_prev = len(self.idx_main)\n",
    "        for i, layer_struct_param in enumerate(struct_param):\n",
    "            num_neurons = layer_struct_param[0]\n",
    "            setattr(self, \"W_{}_main\".format(i), nn.Parameter(torch.randn(num_neurons_prev, num_neurons)))\n",
    "            setattr(self, \"b_{}_main\".format(i), nn.Parameter(torch.zeros(num_neurons)))\n",
    "            init_weight(getattr(self, \"W_{}_main\".format(i)), init=None)\n",
    "            num_neurons_prev = num_neurons\n",
    "            if self.idx_label is not None:\n",
    "                setattr(self, \"W_{}_mul\".format(i), nn.Parameter(torch.randn(len(self.idx_label), num_neurons)))\n",
    "                setattr(self, \"W_{}_add\".format(i), nn.Parameter(torch.randn(len(self.idx_label), num_neurons)))\n",
    "                init_weight(getattr(self, \"W_{}_mul\".format(i)), init=None)\n",
    "                init_weight(getattr(self, \"W_{}_add\".format(i)), init=None)\n",
    "                setattr(self, \"b_{}_mul\".format(i), nn.Parameter(torch.zeros(num_neurons)))\n",
    "                setattr(self, \"b_{}_add\".format(i), nn.Parameter(torch.zeros(num_neurons)))\n",
    "        self.set_cuda(is_cuda)\n",
    "    \n",
    "\n",
    "    def forward(self, input):\n",
    "        output = input[:, self.idx_main]\n",
    "        if self.idx_label is not None:\n",
    "            labels = input[:, self.idx_label]\n",
    "        for i, layer_struct_param in enumerate(self.struct_param):\n",
    "            output = torch.matmul(output, getattr(self, \"W_{}_main\".format(i))) + getattr(self, \"b_{}_main\".format(i))\n",
    "            if \"activation\" in layer_struct_param[2]:\n",
    "                output = get_activation(layer_struct_param[2][\"activation\"])(output)\n",
    "            if self.idx_label is not None:\n",
    "                A_mul = torch.matmul(labels, getattr(self, \"W_{}_mul\".format(i))) + getattr(self, \"b_{}_mul\".format(i))\n",
    "                A_add = torch.matmul(labels, getattr(self, \"W_{}_add\".format(i))) + getattr(self, \"b_{}_add\".format(i))\n",
    "                output = output * A_mul + A_add\n",
    "        return output\n",
    "    \n",
    "    \n",
    "    def get_loss(self, X, y, criterion, **kwargs):\n",
    "        y_pred = self(X)\n",
    "        return criterion(y_pred, y)\n",
    "\n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        if isinstance(is_cuda, str):\n",
    "            self.cuda(is_cuda)\n",
    "        else:\n",
    "            if is_cuda:\n",
    "                self.cuda()\n",
    "            else:\n",
    "                self.cpu()\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        reg = to_Variable([0], is_cuda=self.is_cuda)\n",
    "        return reg\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Labelmix_MLP\"}\n",
    "        model_dict[\"input_size\"] = self.input_size\n",
    "        model_dict[\"struct_param\"] = self.struct_param\n",
    "        if self.idx_label is not None:\n",
    "            model_dict[\"idx_label\"] = to_np_array(self.idx_label).astype(int)\n",
    "        model_dict[\"state_dict\"] = to_cpu_recur(self.state_dict())\n",
    "        return model_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi_MLP (MLPs in series):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Multi_MLP(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        struct_param,\n",
    "        W_init_list = None,     # initialization for weights\n",
    "        b_init_list = None,     # initialization for bias\n",
    "        settings = None,          # Default settings for each layer, if the settings for the layer is not provided in struct_param\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        super(Multi_MLP, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.num_layers = len(struct_param)\n",
    "        self.W_init_list = W_init_list\n",
    "        self.b_init_list = b_init_list\n",
    "        self.settings = deepcopy(settings)\n",
    "        self.num_blocks = len(struct_param)\n",
    "        self.is_cuda = is_cuda\n",
    "        \n",
    "        for i, struct_param_ele in enumerate(struct_param):\n",
    "            input_size_block = input_size if i == 0 else struct_param[i - 1][-1][0]\n",
    "            setattr(self, \"block_{0}\".format(i), MLP(input_size = input_size_block,\n",
    "                                                     struct_param = struct_param_ele,\n",
    "                                                     W_init_list = W_init_list[i] if W_init_list is not None else None,\n",
    "                                                     b_init_list = b_init_list[i] if b_init_list is not None else None,\n",
    "                                                     settings = self.settings[i] if self.settings is not None else {},\n",
    "                                                     is_cuda = self.is_cuda,\n",
    "                                                    ))\n",
    "    \n",
    "    def forward(self, input):\n",
    "        output = input\n",
    "        for i in range(self.num_blocks):\n",
    "            output = getattr(self, \"block_{0}\".format(i))(output)\n",
    "        return output\n",
    "\n",
    "\n",
    "    def get_loss(self, input, target, criterion, **kwargs):\n",
    "        y_pred = self(input, **kwargs)\n",
    "        return criterion(y_pred, target)\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        reg = Variable(torch.FloatTensor([0]), requires_grad = False)\n",
    "        if self.is_cuda:\n",
    "            reg = reg.cuda()\n",
    "        for i in range(self.num_blocks):\n",
    "            reg = reg + getattr(self, \"block_{0}\".format(i)).get_regularization(mode = mode, source = source)\n",
    "        return reg\n",
    "\n",
    "\n",
    "    @property\n",
    "    def struct_param(self):\n",
    "        return [getattr(self, \"block_{0}\".format(i)).struct_param for i in range(self.num_blocks)]\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": self.__class__.__name__}\n",
    "        model_dict[\"input_size\"] = self.input_size\n",
    "        model_dict[\"struct_param\"] = self.struct_param\n",
    "        model_dict[\"weights\"], model_dict[\"bias\"] = self.get_weights_bias(W_source = \"core\", b_source = \"core\")\n",
    "        model_dict[\"settings\"] = deepcopy(self.settings)\n",
    "        model_dict[\"net_type\"] = self.__class__.__name__\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    def load_model_dict(self, model_dict):\n",
    "        new_net = load_model_dict_Multi_MLP(model_dict, is_cuda = self.is_cuda)\n",
    "        self.__dict__.update(new_net.__dict__)\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "\n",
    "    def get_weights_bias(self, W_source = \"core\", b_source = \"core\"):\n",
    "        W_list = []\n",
    "        b_list = []\n",
    "        for i in range(self.num_blocks):\n",
    "            W, b = getattr(self, \"block_{0}\".format(i)).get_weights_bias(W_source = W_source, b_source = b_source)\n",
    "            W_list.append(W)\n",
    "            b_list.append(b)\n",
    "        return deepcopy(W_list), deepcopy(b_list)\n",
    "\n",
    "\n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return {}\n",
    "\n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        for i in range(self.num_blocks):\n",
    "            getattr(self, \"block_{0}\".format(i)).set_cuda(is_cuda)\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        for i in range(self.num_blocks):\n",
    "            getattr(self, \"block_{0}\".format(i)).set_trainable(is_trainable)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Branching_Net:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Branching_Net(nn.Module):\n",
    "    \"\"\"An MLP that consists of a base network, and net_1 and net_2 that branches off from the output of the base network.\"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        net_base_model_dict,\n",
    "        net_1_model_dict,\n",
    "        net_2_model_dict,\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        super(Branching_Net, self).__init__()\n",
    "        self.net_base = load_model_dict(net_base_model_dict, is_cuda = is_cuda)\n",
    "        self.net_1 = load_model_dict(net_1_model_dict, is_cuda = is_cuda)\n",
    "        self.net_2 = load_model_dict(net_2_model_dict, is_cuda = is_cuda)\n",
    "        self.info_dict = {}\n",
    "    \n",
    "    \n",
    "    def forward(self, X, **kwargs):\n",
    "        shared = self.net_base(X)\n",
    "        shared = shared.max(0, keepdim = True)[0]\n",
    "        return self.net_1(shared)[0], self.net_2(shared)[0]\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weights\", \"bias\"], mode = \"L1\"):\n",
    "        reg = self.net_base.get_regularization(source = source, mode = mode) + \\\n",
    "              self.net_1.get_regularization(source = source, mode = mode) + \\\n",
    "              self.net_2.get_regularization(source = source, mode = mode)\n",
    "        return reg\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        self.net_base.set_trainable(is_trainable)\n",
    "        self.net_1.set_trainable(is_trainable)\n",
    "        self.net_2.set_trainable(is_trainable)\n",
    "\n",
    "\n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return deepcopy(self.info_dict)\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Branching_Net\"}\n",
    "        model_dict[\"net_base_model_dict\"] = self.net_base.model_dict\n",
    "        model_dict[\"net_1_model_dict\"] = self.net_1.model_dict\n",
    "        model_dict[\"net_2_model_dict\"] = self.net_2.model_dict\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "    \n",
    "class Fan_in_MLP(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_dict_branch1,\n",
    "        model_dict_branch2,\n",
    "        model_dict_joint,\n",
    "        is_cuda=False,\n",
    "    ):\n",
    "        super(Fan_in_MLP, self).__init__()\n",
    "        if model_dict_branch1 is not None:\n",
    "            self.net_branch1 = load_model_dict(model_dict_branch1, is_cuda=is_cuda)\n",
    "        else:\n",
    "            self.net_branch1 = None\n",
    "        if model_dict_branch2 is not None:\n",
    "            self.net_branch2 = load_model_dict(model_dict_branch2, is_cuda=is_cuda)\n",
    "        else:\n",
    "            self.net_branch2 = None\n",
    "        self.net_joint = load_model_dict(model_dict_joint, is_cuda=is_cuda)\n",
    "        self.is_cuda = is_cuda\n",
    "        self.info_dict = {}\n",
    "    \n",
    "    def forward(self, X1, X2, is_outer=False):\n",
    "        if is_outer:\n",
    "            X2 = X2[...,None,:]\n",
    "        if self.net_branch1 is not None:\n",
    "            X1 = self.net_branch1(X1)\n",
    "        if self.net_branch2 is not None:\n",
    "            X2 = self.net_branch2(X2)\n",
    "        X1, X2 = broadcast_all(X1, X2)\n",
    "        out = torch.cat([X1, X2], -1)\n",
    "        # if is_outer=True, then output dimension: [..., X2dim, X1dim, out_dim]:\n",
    "        return self.net_joint(out).squeeze(-1)\n",
    "    \n",
    "    def get_loss(self, input, target, criterion, **kwargs):\n",
    "        X1, X2 = input\n",
    "        y_pred = self(X1, X2)\n",
    "        return criterion(y_pred, target)\n",
    "    \n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        reg = Variable(torch.FloatTensor([0]), requires_grad = False)\n",
    "        if self.is_cuda:\n",
    "            reg = reg.cuda()\n",
    "        if self.net_branch1 is not None:\n",
    "            reg = reg + self.net_branch1.get_regularization(source=source, mode=mode)\n",
    "        if self.net_branch2 is not None:\n",
    "            reg = reg + self.net_branch2.get_regularization(source=source, mode=mode)\n",
    "        return reg\n",
    "    \n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return deepcopy(self.info_dict)\n",
    "    \n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {'type': self.__class__.__name__}\n",
    "        model_dict[\"model_dict_branch1\"] = self.net_branch1.model_dict if self.net_branch1 is not None else None\n",
    "        model_dict[\"model_dict_branch2\"] = self.net_branch2.model_dict if self.net_branch2 is not None else None\n",
    "        model_dict[\"model_dict_joint\"] = self.net_joint.model_dict\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Mixture_Model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mixture_Model(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_dict_list,\n",
    "        weight_logits_model_dict,\n",
    "        num_components,\n",
    "        is_cuda=False,\n",
    "    ):\n",
    "        super(Mixture_Model, self).__init__()\n",
    "        self.num_components = num_components\n",
    "        for i in range(self.num_components):\n",
    "            if isinstance(model_dict_list, list):\n",
    "                setattr(self, \"model_{}\".format(i), load_model_dict(model_dict_list[i], is_cuda=is_cuda))\n",
    "            else:\n",
    "                assert isinstance(model_dict_list, dict)\n",
    "                setattr(self, \"model_{}\".format(i), load_model_dict(model_dict_list, is_cuda=is_cuda))\n",
    "        self.weight_logits_model = load_model_dict(weight_logits_model_dict, is_cuda=is_cuda)\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "    def forward(self, input):\n",
    "        output_list = []\n",
    "        for i in range(self.num_components):\n",
    "            output = getattr(self, \"model_{}\".format(i))(input)\n",
    "            output_list.append(output)\n",
    "        output_list = torch.stack(output_list, -1)\n",
    "        weight_logits = self.weight_logits_model(input)\n",
    "        return output_list, weight_logits\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Mixture_Model\",\n",
    "                      \"model_dict_list\": [getattr(self, \"model_{}\".format(i)).model_dict for i in range(self.num_components)],\n",
    "                      \"weight_logits_model_dict\": self.weight_logits_model.model_dict,\n",
    "                      \"num_components\": self.num_components,\n",
    "                     }\n",
    "        return model_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model_Ensemble:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model_Ensemble(nn.Module):\n",
    "    \"\"\"Model_Ensemble is a collection of models with the same architecture \n",
    "       but independent parameters\"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_models,\n",
    "        input_size,\n",
    "        struct_param,\n",
    "        W_init_list = None,\n",
    "        b_init_list = None,\n",
    "        settings = None,\n",
    "        net_type = \"MLP\",\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        super(Model_Ensemble, self).__init__()\n",
    "        self.num_models = num_models\n",
    "        self.input_size = input_size\n",
    "        self.net_type = net_type\n",
    "        self.is_cuda = is_cuda\n",
    "        for i in range(self.num_models):\n",
    "            if settings is None:\n",
    "                settings_model = {}\n",
    "            elif isinstance(settings, list) or isinstance(settings, tuple):\n",
    "                settings_model = settings[i]\n",
    "            else:\n",
    "                settings_model = settings\n",
    "            if isinstance(struct_param, tuple):\n",
    "                struct_param_model = struct_param[i]\n",
    "            else:\n",
    "                struct_param_model = struct_param\n",
    "            if net_type == \"MLP\":\n",
    "                net = MLP(input_size = self.input_size,\n",
    "                          struct_param = deepcopy(struct_param_model),\n",
    "                          W_init_list = deepcopy(W_init_list[i]) if W_init_list is not None else None,\n",
    "                          b_init_list = deepcopy(b_init_list[i]) if b_init_list is not None else None,\n",
    "                          settings = deepcopy(settings_model),\n",
    "                          is_cuda = is_cuda,\n",
    "                         )\n",
    "            elif net_type == \"ConvNet\":\n",
    "                net = ConvNet(input_channels = self.input_size,\n",
    "                              struct_param = deepcopy(struct_param_model),\n",
    "                              settings = deepcopy(settings_model),\n",
    "                              is_cuda = is_cuda,\n",
    "                             )\n",
    "            else:\n",
    "                raise Exception(\"Net_type {0} not recognized!\".format(net_type))\n",
    "            setattr(self, \"model_{0}\".format(i), net)\n",
    "\n",
    "\n",
    "    @property\n",
    "    def struct_param(self):\n",
    "        return tuple(getattr(self, \"model_{0}\".format(i)).struct_param for i in range(self.num_models))\n",
    "\n",
    "\n",
    "    @property\n",
    "    def settings(self):\n",
    "        return [getattr(self, \"model_{0}\".format(i)).settings for i in range(self.num_models)]\n",
    "    \n",
    "    \n",
    "    def get_all_models(self):\n",
    "        return [getattr(self, \"model_{0}\".format(i)) for i in range(self.num_models)]\n",
    "\n",
    "\n",
    "    def init_bias_with_input(self, input, mode = \"std_sqrt\", neglect_last_layer = True):\n",
    "        for i in range(self.num_models):\n",
    "            model = getattr(self, \"model_{0}\".format(i))\n",
    "            model.init_bias_with_input(input, mode = mode, neglect_last_layer = neglect_last_layer)\n",
    "    \n",
    "    \n",
    "    def initialize_param_freeze(self, update_values = True):\n",
    "        for i in range(self.num_models):\n",
    "            model = getattr(self, \"model_{0}\".format(i))\n",
    "            model.initialize_param_freeze(update_values = update_values)\n",
    "    \n",
    "    \n",
    "    def apply_model(self, input, model_id):\n",
    "        return fetch_model(self, model_id)(input)\n",
    "\n",
    "\n",
    "    def fetch_model(self, model_id):\n",
    "        return getattr(self, \"model_{0}\".format(model_id))\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        for i in range(self.num_models):\n",
    "            getattr(self, \"model_{0}\".format(i)).set_trainable(is_trainable)\n",
    "\n",
    "\n",
    "    def forward(self, input):\n",
    "        output_list = []\n",
    "        for i in range(self.num_models):\n",
    "            if self.net_type == \"MLP\":\n",
    "                output = getattr(self, \"model_{0}\".format(i))(input)\n",
    "            elif self.net_type == \"ConvNet\":\n",
    "                output = getattr(self, \"model_{0}\".format(i))(input)[0]\n",
    "            else:\n",
    "                raise Exception(\"Net_type {0} not recognized!\".format(self.net_type))\n",
    "            output_list.append(output)\n",
    "        return torch.stack(output_list, 1)\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        if not isinstance(source, list):\n",
    "            source = [source]\n",
    "        reg = Variable(torch.FloatTensor([0]), requires_grad = False)\n",
    "        if self.is_cuda:\n",
    "            reg = reg.cuda()\n",
    "        model0 = self.model_0\n",
    "        # Elastic_weight_reg:\n",
    "        if \"elastic_weight\" in source or \"elastic_bias\" in source:\n",
    "            # Setting up excluded layer:\n",
    "            excluded_layer = kwargs[\"excluded_layer\"] if \"excluded_layer\" in kwargs else [-1]\n",
    "            if not isinstance(excluded_layer, list):\n",
    "                excluded_layer = [excluded_layer]\n",
    "            excluded_layer = [element + model0.num_layers if element < 0 else element for element in excluded_layer]\n",
    "            elastic_mode = kwargs[\"elastic_mode\"] if \"elastic_mode\" in kwargs else \"var\"\n",
    "            \n",
    "            # Compute the elastic_weight_reg:\n",
    "            for k in range(model0.num_layers):\n",
    "                if k in excluded_layer:\n",
    "                    continue\n",
    "                W_accum_k = []\n",
    "                b_accum_k = []\n",
    "                num_neurons_prev = model0.struct_param[k - 1][0] if k > 0 else self.input_size\n",
    "                num_neurons = model0.struct_param[k][0]\n",
    "                for i in range(self.num_models):\n",
    "                    model = getattr(self, \"model_{0}\".format(i))\n",
    "                    assert model0.num_layers == model.num_layers\n",
    "                    assert num_neurons_prev == model.struct_param[k - 1][0] if k > 0 else model.input_size, \\\n",
    "                            \"all models' input/output size at each layer must be identical!\"\n",
    "                    assert num_neurons == model.struct_param[k][0], \\\n",
    "                            \"all models' input/output size at each layer must be identical!\"\n",
    "                    layer_k = getattr(model, \"layer_{0}\".format(k))\n",
    "                    if \"elastic_weight\" in source:\n",
    "                        W_accum_k.append(layer_k.W_core)\n",
    "                    if \"elastic_bias\" in source:\n",
    "                        b_accum_k.append(layer_k.b_core)\n",
    "                if \"elastic_weight\" in source:\n",
    "                    if elastic_mode == \"var\":\n",
    "                        reg = reg + torch.stack(W_accum_k, -1).var(-1).sum()\n",
    "                    elif elastic_mode == \"std\":\n",
    "                        reg = reg + torch.stack(W_accum_k, -1).std(-1).sum()\n",
    "                    else:\n",
    "                        raise\n",
    "                if \"elastic_bias\" in source:\n",
    "                    if elastic_mode == \"var\":\n",
    "                        reg = reg + torch.stack(b_accum_k, -1).var(-1).sum()\n",
    "                    elif elastic_mode == \"std\":\n",
    "                        reg = reg + torch.stack(b_accum_k, -1).std(-1).sum()\n",
    "                    else:\n",
    "                        raise\n",
    "            source_core = deepcopy(source)\n",
    "            if \"elastic_weight\" in source_core:\n",
    "                source_core.remove(\"elastic_weight\")\n",
    "            if \"elastic_bias\" in source_core:\n",
    "                source_core.remove(\"elastic_bias\")\n",
    "        else:\n",
    "            source_core = source\n",
    "        \n",
    "        # Other regularizations:\n",
    "        for k in range(self.num_models):\n",
    "            reg = reg + getattr(self, \"model_{0}\".format(k)).get_regularization(source = source_core, mode = mode, **kwargs)\n",
    "        return reg\n",
    "    \n",
    "    \n",
    "    def get_weights_bias(self, W_source = None, b_source = None, verbose = False, isplot = False):\n",
    "        W_list_dict = {}\n",
    "        b_list_dict = {}\n",
    "        for i in range(self.num_models):\n",
    "            if verbose:\n",
    "                print(\"\\nmodel {0}:\".format(i))\n",
    "            W_list_dict[i], b_list_dict[i] = getattr(self, \"model_{0}\".format(i)).get_weights_bias(\n",
    "                W_source = W_source, b_source = b_source, verbose = verbose, isplot = isplot)\n",
    "        return W_list_dict, b_list_dict\n",
    "    \n",
    "    \n",
    "    def combine_to_net(self, mode = \"mean\", last_layer_mode = \"concatenate\"):\n",
    "        model0 = self.model_0\n",
    "        if mode == \"mean\":\n",
    "            struct_param = deepcopy(model0.struct_param)\n",
    "            settings = deepcopy(model0.settings)\n",
    "            W_init_list = []\n",
    "            b_init_list = []\n",
    "            for k in range(model0.num_layers):\n",
    "                num_neurons_prev = model0.struct_param[k - 1][0] if k > 0 else self.input_size\n",
    "                num_neurons = model0.struct_param[k][0]\n",
    "                W_accum_k = []\n",
    "                b_accum_k = []\n",
    "                for i in range(self.num_models):\n",
    "                    model = getattr(self, \"model_{0}\".format(i))\n",
    "                    assert model0.num_layers == model.num_layers\n",
    "                    assert num_neurons_prev == model.struct_param[k - 1][0] if k > 0 else model.input_size, \\\n",
    "                            \"If mode == 'mean', all models' input/output size at each layer must be identical!\"\n",
    "                    assert num_neurons == model.struct_param[k][0], \\\n",
    "                            \"If mode == 'mean', all models' input/output size at each layer must be identical!\"\n",
    "                    layer_k = getattr(model, \"layer_{0}\".format(k))\n",
    "                    W_accum_k.append(layer_k.W_core)\n",
    "                    b_accum_k.append(layer_k.b_core)\n",
    "\n",
    "                if k == model0.num_layers - 1:\n",
    "                    current_mode = last_layer_mode\n",
    "                else:\n",
    "                    current_mode = mode\n",
    "\n",
    "                if current_mode == \"mean\":\n",
    "                    W_accum_k = torch.stack(W_accum_k, -1).mean(-1)\n",
    "                    b_accum_k = torch.stack(b_accum_k, -1).mean(-1)\n",
    "                elif current_mode == \"concatenate\":\n",
    "                    W_accum_k = torch.cat(W_accum_k, -1)\n",
    "                    b_accum_k = torch.cat(b_accum_k, -1)\n",
    "                    struct_param[-1][0] = sum([self.struct_param[i][-1][0] for i in range(self.num_models)])\n",
    "                else:\n",
    "                    raise Exception(\"mode {0} not recognized!\".format(last_layer_mode))\n",
    "                W_init_list.append(W_accum_k.data.numpy())\n",
    "                b_init_list.append(b_accum_k.data.numpy())\n",
    "            \n",
    "            # Build the net:\n",
    "            net = MLP(input_size = self.input_size,\n",
    "                      struct_param = struct_param,\n",
    "                      W_init_list = W_init_list, \n",
    "                      b_init_list = b_init_list,\n",
    "                      settings = settings,\n",
    "                     )\n",
    "        else:\n",
    "            raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "        return net\n",
    "        \n",
    "    \n",
    "    def remove_models(self, model_ids):\n",
    "        if not isinstance(model_ids, list):\n",
    "            model_ids = [model_ids]\n",
    "        model_list = []\n",
    "        k = 0\n",
    "        for i in range(self.num_models):\n",
    "            if i not in model_ids:\n",
    "                if k != i:\n",
    "                    setattr(self, \"model_{0}\".format(k), getattr(self, \"model_{0}\".format(i)))\n",
    "                k += 1\n",
    "        num_models_new = k\n",
    "        for i in range(num_models_new, self.num_models):\n",
    "            delattr(self, \"model_{0}\".format(i))\n",
    "        self.num_models = num_models_new\n",
    "\n",
    "\n",
    "    def add_models(self, models):\n",
    "        if not isinstance(models, list):\n",
    "            models = [models]\n",
    "        for i, model in enumerate(models):\n",
    "            setattr(self, \"model_{0}\".format(i + self.num_models), model)\n",
    "        self.num_models += len(models)\n",
    "\n",
    "\n",
    "    def simplify(self, X, y, idx, mode = \"full\", validation_data = None, isplot = False, **kwargs):\n",
    "        def process_idx(idx):\n",
    "            idx = idx.byte()\n",
    "            if len(idx.size()) == 1:\n",
    "                idx = idx.unqueeze(1)\n",
    "            if idx.size(1) == 1:\n",
    "                idx = idx.repeat(1, self.num_models)\n",
    "            return idx\n",
    "        idx = process_idx(idx)\n",
    "        if validation_data is not None:\n",
    "            X_valid, y_valid, idx_valid = validation_data\n",
    "            idx_valid = process_idx(idx_valid)        \n",
    "        \n",
    "        loss_dict = {}\n",
    "        for i in range(self.num_models):\n",
    "            model = getattr(self, \"model_{0}\".format(i))\n",
    "            X_chosen = torch.masked_select(X, idx[:, i:i+1]).view(-1, X.size(1))\n",
    "            y_chosen = torch.masked_select(y, idx[:, i:i+1]).view(-1, y.size(1))\n",
    "            if validation_data is not None:\n",
    "                X_valid_chosen = torch.masked_select(X_valid, idx_valid[:, i:i+1]).view(-1, X_valid.size(1))\n",
    "                y_valid_chosen = torch.masked_select(y_valid, idx_valid[:, i:i+1]).view(-1, y_valid.size(1))\n",
    "                if len(X_valid_chosen) == 0:\n",
    "                    validation_data_chosen = None\n",
    "                else:\n",
    "                    validation_data_chosen = (X_valid_chosen, y_valid_chosen)\n",
    "            else:\n",
    "                validation_data_chosen = None\n",
    "            if len(X_chosen) == 0:\n",
    "                print(\"The {0}'th model has no corresponding data to simplify with, skip.\".format(i))\n",
    "            else:\n",
    "                new_model, loss_dict[\"model_{0}\".format(i)] = simplify(model, X_chosen, y_chosen, mode = mode, validation_data = validation_data_chosen, isplot = isplot, target_name = \"model_{0}\".format(i), **kwargs)\n",
    "                setattr(self, \"model_{0}\".format(i), new_model)\n",
    "        return loss_dict\n",
    "    \n",
    "    \n",
    "    def get_sympy_expression(self):\n",
    "        expressions = {}\n",
    "        for k in range(self.num_models):\n",
    "            print(\"\\nmodel {0}:\".format(k))\n",
    "            expressions[\"model_{0}\".format(k)] = getattr(self, \"model_{0}\".format(k)).get_sympy_expression()\n",
    "        return expressions\n",
    "\n",
    "\n",
    "    @property\n",
    "    def DL(self):\n",
    "        return np.sum([getattr(self, \"model_{0}\".format(i)).DL for i in range(self.num_models)])\n",
    "\n",
    "\n",
    "    def get_weights_bias(self, W_source = None, b_source = None, verbose = False, isplot = False):\n",
    "        W_list_dict = {}\n",
    "        b_list_dict = {}\n",
    "        for i in range(self.num_models):\n",
    "            if verbose:\n",
    "                print(\"\\nmodel {0}:\".format(i))\n",
    "            W_list_dict[i], b_list_dict[i] = getattr(self, \"model_{0}\".format(i)).get_weights_bias(W_source = W_source, b_source = b_source, verbose = verbose, isplot = isplot)\n",
    "        return W_list_dict, b_list_dict\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Model_Ensemble\"}\n",
    "        for i in range(self.num_models):\n",
    "            model_dict[\"model_{0}\".format(i)] = getattr(self, \"model_{0}\".format(i)).model_dict\n",
    "        model_dict[\"input_size\"] = self.input_size\n",
    "        model_dict[\"struct_param\"] = self.struct_param\n",
    "        model_dict[\"num_models\"] = self.num_models\n",
    "        model_dict[\"net_type\"] = self.net_type\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    def load_model_dict(self, model_dict):\n",
    "        new_model_ensemble = load_model_dict(model_dict, is_cuda = self.is_cuda)\n",
    "        self.__dict__.update(new_model_ensemble.__dict__)\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "        \n",
    "def load_model_dict_model_ensemble(model_dict, is_cuda = False):\n",
    "    num_models = len([model_name for model_name in model_dict if model_name[:6] == \"model_\"])\n",
    "    return Model_Ensemble(num_models = num_models,\n",
    "                          input_size = model_dict[\"input_size\"],\n",
    "                          struct_param = tuple([deepcopy(model_dict[\"model_{0}\".format(i)][\"struct_param\"]) for i in range(num_models)]),\n",
    "                          W_init_list = [deepcopy(model_dict[\"model_{0}\".format(i)][\"weights\"]) for i in range(num_models)],\n",
    "                          b_init_list = [deepcopy(model_dict[\"model_{0}\".format(i)][\"bias\"]) for i in range(num_models)],\n",
    "                          settings = [deepcopy(model_dict[\"model_{0}\".format(i)][\"settings\"]) for i in range(num_models)],\n",
    "                          net_type = model_dict[\"net_type\"] if \"net_type\" in model_dict else \"MLP\",\n",
    "                          is_cuda = is_cuda,\n",
    "                         )\n",
    "\n",
    "\n",
    "def combine_model_ensembles(model_ensembles, input_size):\n",
    "    model_ensembles = deepcopy(model_ensembles)\n",
    "    model_ensemble_combined = None\n",
    "    model_id = 0\n",
    "    for k, model_ensemble in enumerate(model_ensembles):\n",
    "        if model_ensemble.input_size == input_size:\n",
    "            if model_ensemble_combined is None:\n",
    "                model_ensemble_combined = model_ensemble\n",
    "        else:\n",
    "            continue  \n",
    "        for i in range(model_ensemble.num_models):\n",
    "            model = getattr(model_ensemble, \"model_{0}\".format(i))\n",
    "            setattr(model_ensemble_combined, \"model_{0}\".format(model_id), model)\n",
    "            model_id += 1\n",
    "    model_ensemble_combined.num_models = model_id\n",
    "    return model_ensemble_combined\n",
    "\n",
    "\n",
    "def construct_model_ensemble_from_nets(nets):\n",
    "    num_models = len(nets)\n",
    "    if num_models is None:\n",
    "        return None\n",
    "    input_size = nets[0].input_size\n",
    "    struct_param = tuple(net.struct_param for net in nets)\n",
    "    is_cuda = False\n",
    "    for net in nets:\n",
    "        if net.input_size != input_size:\n",
    "            raise Exception(\"The input_size for all nets must be the same!\")\n",
    "        if net.is_cuda:\n",
    "            is_cuda = True\n",
    "    model_ensemble = Model_Ensemble(num_models = num_models, input_size = input_size, struct_param = struct_param, is_cuda = is_cuda)\n",
    "    for i, net in enumerate(nets):\n",
    "        setattr(model_ensemble, \"model_{0}\".format(i), net)\n",
    "    return model_ensemble"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model_with_uncertainty(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_pred,\n",
    "        model_logstd,\n",
    "        ):\n",
    "        super(Model_with_uncertainty, self).__init__()\n",
    "        self.model_pred = model_pred\n",
    "        self.model_logstd = model_logstd\n",
    "        \n",
    "    def forward(self, input, noise_amp = None, **kwargs):\n",
    "        return self.model_pred(input, noise_amp = noise_amp, **kwargs), self.model_logstd(input, **kwargs)\n",
    "    \n",
    "    def get_loss(self, input, target, criterion, noise_amp = None, **kwargs):\n",
    "        pred, log_std = self(input, noise_amp = noise_amp, **kwargs)\n",
    "        return criterion(pred = pred, target = target, log_std = log_std)\n",
    "    \n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        return self.model_pred.get_regularization(source = source, mode = mode, **kwargs) + self.model_logstd.get_regularization(source = source, mode = mode, **kwargs)\n",
    "    \n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {}\n",
    "        model_dict[\"type\"] = \"Model_with_Uncertainty\"\n",
    "        model_dict[\"model_pred\"] = self.model_pred.model_dict\n",
    "        model_dict[\"model_logstd\"] = self.model_logstd.model_dict\n",
    "        return model_dict\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        self.model_pred.set_cuda(is_cuda)\n",
    "        self.model_logstd.set_cuda(is_cuda)\n",
    "        \n",
    "    def set_trainable(self, is_trainable):\n",
    "        self.model_pred.set_trainable(is_trainable)\n",
    "        self.model_logstd.set_trainable(is_trainable)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RNN:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class RNNCellBase(nn.Module):\n",
    "    def extra_repr(self):\n",
    "        s = '{input_size}, {hidden_size}'\n",
    "        if 'bias' in self.__dict__ and self.bias is not True:\n",
    "            s += ', bias={bias}'\n",
    "        if 'nonlinearity' in self.__dict__ and self.nonlinearity != \"tanh\":\n",
    "            s += ', nonlinearity={nonlinearity}'\n",
    "        return s.format(**self.__dict__)\n",
    "\n",
    "    def check_forward_input(self, input):\n",
    "        if input.size(1) != self.input_size:\n",
    "            raise RuntimeError(\n",
    "                \"input has inconsistent input_size: got {}, expected {}\".format(\n",
    "                    input.size(1), self.input_size))\n",
    "\n",
    "    def check_forward_hidden(self, input, hx, hidden_label=''):\n",
    "        if input.size(0) != hx.size(0):\n",
    "            raise RuntimeError(\n",
    "                \"Input batch size {} doesn't match hidden{} batch size {}\".format(\n",
    "                    input.size(0), hidden_label, hx.size(0)))\n",
    "\n",
    "        if hx.size(1) != self.hidden_size:\n",
    "            raise RuntimeError(\n",
    "                \"hidden{} has inconsistent hidden_size: got {}, expected {}\".format(\n",
    "                    hidden_label, hx.size(1), self.hidden_size))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LSTM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LSTM(RNNCellBase):\n",
    "    \"\"\"a LSTM class\"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_size,\n",
    "        hidden_size,\n",
    "        output_struct_param,\n",
    "        output_settings = {},\n",
    "        bias = True,\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        super(LSTM, self).__init__()\n",
    "        self.input_size = input_size\n",
    "        self.hidden_size = hidden_size\n",
    "        self.bias = bias\n",
    "        self.W_ih = nn.Parameter(torch.Tensor(4 * hidden_size, input_size))\n",
    "        self.W_hh = nn.Parameter(torch.Tensor(4 * hidden_size, hidden_size))\n",
    "        self.output_net = MLP(input_size = self.hidden_size, struct_param = output_struct_param, settings = output_settings, is_cuda = is_cuda)\n",
    "        if bias:\n",
    "            self.b_ih = nn.Parameter(torch.Tensor(4 * hidden_size))\n",
    "            self.b_hh = nn.Parameter(torch.Tensor(4 * hidden_size))\n",
    "        else:\n",
    "            self.register_parameter('b_ih', None)\n",
    "            self.register_parameter('b_hh', None)\n",
    "        self.reset_parameters()\n",
    "        self.is_cuda = is_cuda\n",
    "        self.device = torch.device(self.is_cuda if isinstance(self.is_cuda, str) else \"cuda\" if self.is_cuda else \"cpu\")\n",
    "        self.to(self.device)\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        stdv = 1.0 / np.sqrt(self.hidden_size)\n",
    "        for weight in self.parameters():\n",
    "            weight.data.uniform_(-stdv, stdv)\n",
    "\n",
    "    def forward_one_step(self, input, hx):\n",
    "        self.check_forward_input(input)\n",
    "        self.check_forward_hidden(input, hx[0], '[0]')\n",
    "        self.check_forward_hidden(input, hx[1], '[1]')\n",
    "        return self._backend.LSTMCell(\n",
    "            input, hx,\n",
    "            self.W_ih, self.W_hh,\n",
    "            self.b_ih, self.b_hh,\n",
    "        )\n",
    "    \n",
    "    def forward(self, input, hx = None):\n",
    "        if hx is None:\n",
    "            hx = [torch.randn(input.size(0), self.hidden_size).to(self.device),\n",
    "                  torch.randn(input.size(0), self.hidden_size).to(self.device),\n",
    "                 ]\n",
    "        hhx, ccx = hx\n",
    "        for i in range(input.size(1)):\n",
    "            hhx, ccx = self.forward_one_step(input[:, i], (hhx, ccx))\n",
    "        output = self.output_net(hhx)\n",
    "        return output\n",
    "\n",
    "    def get_regularization(self, source, mode = \"L1\", **kwargs):\n",
    "        if not isinstance(source, list):\n",
    "            source = [source]\n",
    "        reg = self.output_net.get_regularization(source = source, mode = mode)\n",
    "        for source_ele in source:\n",
    "            if source_ele == \"weight\":\n",
    "                if mode == \"L1\":\n",
    "                    reg = reg + self.W_ih.abs().sum() + self.W_hh.abs().sum()\n",
    "                elif mode == \"L2\":\n",
    "                    reg = reg + (self.W_ih ** 2).sum() + (self.W_hh ** 2).sum()\n",
    "                else:\n",
    "                    raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "            elif source_ele == \"bias\":\n",
    "                if self.bias:\n",
    "                    if mode == \"L1\":\n",
    "                        reg = reg + self.b_ih.abs().sum() + self.b_hh.abs().sum()\n",
    "                    elif mode == \"L2\":\n",
    "                        reg = reg + (self.b_ih ** 2).sum() + (self.b_hh ** 2).sum()\n",
    "                    else:\n",
    "                        raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "            else:\n",
    "                raise Exception(\"source {0} not recognized!\".format(source_ele))\n",
    "        return reg\n",
    "    \n",
    "    def get_weights_bias(self, W_source = None, b_source = None, verbose = False, isplot = False):\n",
    "        W_dict = OrderedDict()\n",
    "        b_dict = OrderedDict()\n",
    "        W_o, b_o = self.output_net.get_weights_bias(W_source = W_source, b_source = b_source)\n",
    "        if W_source == \"core\":\n",
    "            W_dict[\"W_ih\"] = self.W_ih.cpu().detach().numpy()\n",
    "            W_dict[\"W_hh\"] = self.W_hh.cpu().detach().numpy()\n",
    "            W_dict[\"W_o\"] = W_o\n",
    "            if isplot:\n",
    "                print(\"W_ih, W_hh:\")\n",
    "                plot_matrices([W_dict[\"W_ih\"], W_dict[\"W_hh\"]])\n",
    "                print(\"W_o:\")\n",
    "                plot_matrices(W_o)\n",
    "        if self.bias and b_source == \"core\":\n",
    "            b_dict[\"b_ih\"] = self.b_ih.cpu().detach().numpy()\n",
    "            b_dict[\"b_hh\"] = self.b_hh.cpu().detach().numpy()\n",
    "            b_dict[\"b_o\"] = b_o\n",
    "            if isplot:\n",
    "                print(\"b_ih, b_hh:\")\n",
    "                plot_matrices([b_dict[\"b_ih\"], b_dict[\"b_hh\"]])\n",
    "                print(\"b_o:\")\n",
    "                plot_matrices(b_o)\n",
    "        return W_dict, b_dict\n",
    "    \n",
    "    def get_loss(self, input, target, criterion, hx = None, **kwargs):\n",
    "        y_pred = self(input, hx = hx)\n",
    "        return criterion(y_pred, target)\n",
    "    \n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return {}\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wide ResNet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def conv3x3(in_planes, out_planes, stride=1):\n",
    "    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=True)\n",
    "\n",
    "def conv_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        init.xavier_uniform_(m.weight, gain=np.sqrt(2))\n",
    "        init.constant_(m.bias, 0)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        init.constant_(m.weight, 1)\n",
    "        init.constant_(m.bias, 0)\n",
    "\n",
    "\n",
    "class wide_basic(nn.Module):\n",
    "    def __init__(self, in_planes, planes, dropout_rate=None, stride=1):\n",
    "        super(wide_basic, self).__init__()\n",
    "        self.bn1 = nn.BatchNorm2d(in_planes)\n",
    "        self.conv1 = nn.Conv2d(in_planes, planes, kernel_size=3, padding=1, bias=True)\n",
    "        if dropout_rate is not None:\n",
    "            self.dropout = nn.Dropout(p=dropout_rate)\n",
    "        else:\n",
    "            self.dropout = None\n",
    "        self.bn2 = nn.BatchNorm2d(planes)\n",
    "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, padding=1, bias=True)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if stride != 1 or in_planes != planes:\n",
    "            self.shortcut = nn.Sequential(\n",
    "                nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride, bias=True),\n",
    "            )\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv1(F.relu(self.bn1(x)))\n",
    "        if self.dropout is not None:\n",
    "            out = self.dropout(out)\n",
    "        out = self.conv2(F.relu(self.bn2(out)))\n",
    "        out += self.shortcut(x)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Wide_ResNet(nn.Module):\n",
    "    \"\"\"Adapted from https://github.com/meliketoy/wide-resnet.pytorch/blob/master/networks/wide_resnet.py\"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        depth,\n",
    "        widen_factor,\n",
    "        input_channels,\n",
    "        output_size,\n",
    "        dropout_rate=None,\n",
    "        is_cuda=False,\n",
    "    ):\n",
    "        super(Wide_ResNet, self).__init__()\n",
    "\n",
    "        self.depth = depth\n",
    "        self.widen_factor = widen_factor\n",
    "        self.input_channels = input_channels\n",
    "        self.dropout_rate = dropout_rate\n",
    "        self.output_size = output_size\n",
    "\n",
    "        assert ((depth-4)%6 ==0), 'Wide-resnet depth should be 6n+4'\n",
    "        n = (depth-4)//6\n",
    "        k = widen_factor\n",
    "\n",
    "        nStages = [16*k, 16*k, 32*k, 64*k]\n",
    "        self.in_planes = nStages[0]\n",
    "\n",
    "        self.conv1 = conv3x3(self.input_channels,nStages[0])\n",
    "        self.layer1 = self._wide_layer(wide_basic, nStages[1], n, dropout_rate, stride=1)\n",
    "        self.layer2 = self._wide_layer(wide_basic, nStages[2], n, dropout_rate, stride=2)\n",
    "        self.layer3 = self._wide_layer(wide_basic, nStages[3], n, dropout_rate, stride=2)\n",
    "        self.bn1 = nn.BatchNorm2d(nStages[3], momentum=0.9)\n",
    "        self.linear = nn.Linear(nStages[3], output_size)\n",
    "        self.set_cuda(is_cuda)\n",
    "\n",
    "    def _wide_layer(self, block, planes, num_blocks, dropout_rate, stride):\n",
    "        strides = [stride] + [1]*(int(num_blocks)-1)\n",
    "        layers = []\n",
    "\n",
    "        for stride in strides:\n",
    "            layers.append(block(self.in_planes, planes, dropout_rate, stride))\n",
    "            self.in_planes = planes\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv1(x)\n",
    "        out = self.layer1(out)\n",
    "        out = self.layer2(out)\n",
    "        out = self.layer3(out)\n",
    "        out = F.relu(self.bn1(out))\n",
    "        out = out.mean((-1,-2))  # replacing the out= F.avg_pool2d(out, 8) which is sensitive to the input shape.\n",
    "        out = out.view(out.size(0), -1)\n",
    "        out = self.linear(out)\n",
    "\n",
    "        return out\n",
    "    \n",
    "    def set_cuda(self, is_cuda):\n",
    "        if isinstance(is_cuda, str):\n",
    "            self.cuda(is_cuda)\n",
    "        else:\n",
    "            if is_cuda:\n",
    "                self.cuda()\n",
    "            else:\n",
    "                self.cpu()\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "    \n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Wide_ResNet\"}\n",
    "        model_dict[\"state_dict\"] = to_cpu_recur(self.state_dict())\n",
    "        model_dict[\"depth\"] = self.depth\n",
    "        model_dict[\"widen_factor\"] = self.widen_factor\n",
    "        model_dict[\"input_channels\"] = self.input_channels\n",
    "        model_dict[\"output_size\"] = self.output_size\n",
    "        model_dict[\"dropout_rate\"] = self.dropout_rate\n",
    "        return model_dict\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "    \n",
    "    def get_regularization(self, *args, **kwargs):\n",
    "        return to_Variable([0], is_cuda = self.is_cuda)\n",
    "    \n",
    "    def prepare_inspection(self, *args, **kwargs):\n",
    "        return {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CNN:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ConvNet(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_channels,\n",
    "        struct_param=None,\n",
    "        W_init_list=None,\n",
    "        b_init_list=None,\n",
    "        settings={},\n",
    "        return_indices=False,\n",
    "        is_cuda=False,\n",
    "        ):\n",
    "        super(ConvNet, self).__init__()\n",
    "        self.input_channels = input_channels\n",
    "        if struct_param is not None:\n",
    "            self.struct_param = struct_param\n",
    "            self.W_init_list = W_init_list\n",
    "            self.b_init_list = b_init_list\n",
    "            self.settings = settings\n",
    "            self.num_layers = len(struct_param)\n",
    "            self.info_dict = {}\n",
    "            self.param_available = [\"Conv2d\", \"ConvTranspose2d\", \"BatchNorm2d\", \"Simple_Layer\"]\n",
    "            self.return_indices = return_indices\n",
    "            for i in range(len(self.struct_param)):\n",
    "                if i > 0:\n",
    "                    k = 1\n",
    "                    while self.struct_param[i - k][0] is None:\n",
    "                        k += 1\n",
    "                    num_channels_prev = self.struct_param[i - k][0]\n",
    "                else:\n",
    "                    num_channels_prev = input_channels\n",
    "                    k = 0\n",
    "                if self.struct_param[i - k][1] == \"Simple_Layer\" and isinstance(num_channels_prev, tuple) and len(num_channels_prev) == 3:\n",
    "                    num_channels_prev = num_channels_prev[0]\n",
    "                num_channels = self.struct_param[i][0]\n",
    "                layer_type = self.struct_param[i][1]\n",
    "                layer_settings = self.struct_param[i][2]\n",
    "                if \"layer_input_size\" in layer_settings and isinstance(layer_settings[\"layer_input_size\"], tuple):\n",
    "                    num_channels_prev = layer_settings[\"layer_input_size\"][0]\n",
    "                if layer_type == \"Conv2d\":\n",
    "                    layer = nn.Conv2d(num_channels_prev, \n",
    "                                      num_channels,\n",
    "                                      kernel_size = layer_settings[\"kernel_size\"],\n",
    "                                      stride = layer_settings[\"stride\"] if \"stride\" in layer_settings else 1,\n",
    "                                      padding = layer_settings[\"padding\"] if \"padding\" in layer_settings else 0,\n",
    "                                      dilation = layer_settings[\"dilation\"] if \"dilation\" in layer_settings else 1,\n",
    "                                     )\n",
    "                elif layer_type == \"ConvTranspose2d\":\n",
    "                    layer = nn.ConvTranspose2d(num_channels_prev,\n",
    "                                               num_channels,\n",
    "                                               kernel_size = layer_settings[\"kernel_size\"],\n",
    "                                               stride = layer_settings[\"stride\"] if \"stride\" in layer_settings else 1,\n",
    "                                               padding = layer_settings[\"padding\"] if \"padding\" in layer_settings else 0,\n",
    "                                               output_padding = layer_settings[\"output_padding\"] if \"output_padding\" in layer_settings else 0,\n",
    "                                               dilation = layer_settings[\"dilation\"] if \"dilation\" in layer_settings else 1,\n",
    "                                              )\n",
    "                elif layer_type == \"Simple_Layer\":\n",
    "                    layer = get_Layer(layer_type = layer_type,\n",
    "                                      input_size = layer_settings[\"layer_input_size\"],\n",
    "                                      output_size = num_channels,\n",
    "                                      W_init = W_init_list[i] if self.W_init_list is not None and self.W_init_list[i] is not None else None,\n",
    "                                      b_init = b_init_list[i] if self.b_init_list is not None and self.b_init_list[i] is not None else None,\n",
    "                                      settings = layer_settings,\n",
    "                                      is_cuda = is_cuda,\n",
    "                                     )\n",
    "                elif layer_type == \"MaxPool2d\":\n",
    "                    layer = nn.MaxPool2d(kernel_size = layer_settings[\"kernel_size\"],\n",
    "                                         stride = layer_settings[\"stride\"] if \"stride\" in layer_settings else None,\n",
    "                                         padding = layer_settings[\"padding\"] if \"padding\" in layer_settings else 0,\n",
    "                                         return_indices = layer_settings[\"return_indices\"] if \"return_indices\" in layer_settings else False,\n",
    "                                        )\n",
    "                elif layer_type == \"MaxUnpool2d\":\n",
    "                    layer = nn.MaxUnpool2d(kernel_size = layer_settings[\"kernel_size\"],\n",
    "                                           stride = layer_settings[\"stride\"] if \"stride\" in layer_settings else None,\n",
    "                                           padding = layer_settings[\"padding\"] if \"padding\" in layer_settings else 0,\n",
    "                                          )\n",
    "                elif layer_type == \"Upsample\":\n",
    "                    layer = nn.Upsample(scale_factor = layer_settings[\"scale_factor\"],\n",
    "                                        mode = layer_settings[\"mode\"] if \"mode\" in layer_settings else \"nearest\",\n",
    "                                       )\n",
    "                elif layer_type == \"BatchNorm2d\":\n",
    "                    layer = nn.BatchNorm2d(num_features = num_channels)\n",
    "                elif layer_type == \"Dropout2d\":\n",
    "                    layer = nn.Dropout2d(p = 0.5)\n",
    "                elif layer_type == \"Flatten\":\n",
    "                    layer = Flatten()\n",
    "                else:\n",
    "                    raise Exception(\"layer_type {0} not recognized!\".format(layer_type))\n",
    "\n",
    "                # Initialize using provided initial values:\n",
    "                if self.W_init_list is not None and self.W_init_list[i] is not None and layer_type not in [\"Simple_Layer\"]:\n",
    "                    layer.weight.data = torch.FloatTensor(self.W_init_list[i])\n",
    "                    layer.bias.data = torch.FloatTensor(self.b_init_list[i])\n",
    "\n",
    "                setattr(self, \"layer_{0}\".format(i), layer)\n",
    "            self.set_cuda(is_cuda)\n",
    "\n",
    "\n",
    "    def forward(self, input, indices_list = None, **kwargs):\n",
    "        return self.inspect_operation(input, operation_between = (0, self.num_layers), indices_list = indices_list)\n",
    "    \n",
    "    \n",
    "    def inspect_operation(self, input, operation_between, indices_list = None):\n",
    "        output = input\n",
    "        if indices_list is None:\n",
    "            indices_list = []\n",
    "        start_layer, end_layer = operation_between\n",
    "        if end_layer < 0:\n",
    "            end_layer += self.num_layers\n",
    "        for i in range(start_layer, end_layer):\n",
    "            if \"layer_input_size\" in self.struct_param[i][2]:\n",
    "                output_size_last = output.shape[0]\n",
    "                layer_input_size = self.struct_param[i][2][\"layer_input_size\"]\n",
    "                if not isinstance(layer_input_size, tuple):\n",
    "                    layer_input_size = (layer_input_size,)\n",
    "                output = output.view(-1, *layer_input_size)\n",
    "                assert output.shape[0] == output_size_last, \"output_size reshaped to different length. Check shape!\"\n",
    "            if \"Unpool\" in self.struct_param[i][1]:\n",
    "                output_tentative = getattr(self, \"layer_{0}\".format(i))(output, indices_list.pop(-1))\n",
    "            else:\n",
    "                output_tentative = getattr(self, \"layer_{0}\".format(i))(output)\n",
    "            if isinstance(output_tentative, tuple):\n",
    "                output, indices = output_tentative\n",
    "                indices_list.append(indices)\n",
    "            else:\n",
    "                output = output_tentative\n",
    "            if \"activation\" in self.struct_param[i][2]:\n",
    "                activation = self.struct_param[i][2][\"activation\"]\n",
    "            else:\n",
    "                if \"activation\" in self.settings:\n",
    "                    activation = self.settings[\"activation\"]\n",
    "                else:\n",
    "                    activation = \"linear\"\n",
    "                if \"Pool\" in self.struct_param[i][1] or \"Unpool\" in self.struct_param[i][1] or \"Upsample\" in self.struct_param[i][1]:\n",
    "                    activation = \"linear\"\n",
    "            output = get_activation(activation)(output)\n",
    "        if self.return_indices:\n",
    "            return output, indices_list\n",
    "        else:\n",
    "            return output\n",
    "\n",
    "\n",
    "    def get_loss(self, input, target, criterion, **kwargs):\n",
    "        y_pred = self(input, **kwargs)\n",
    "        if self.return_indices:\n",
    "            y_pred = y_pred[0]\n",
    "        return criterion(y_pred, target)\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        if not isinstance(source, list):\n",
    "            source = [source]\n",
    "        reg = Variable(torch.FloatTensor([0]), requires_grad = False)\n",
    "        if self.is_cuda:\n",
    "            reg = reg.cuda()\n",
    "        for k in range(self.num_layers):\n",
    "            if self.struct_param[k][1] not in self.param_available:\n",
    "                continue\n",
    "            layer = getattr(self, \"layer_{0}\".format(k))\n",
    "            for source_ele in source:\n",
    "                if source_ele == \"weight\":\n",
    "                    if self.struct_param[k][1] not in [\"Simple_Layer\"]:\n",
    "                        item = layer.weight\n",
    "                    else:\n",
    "                        item = layer.W_core\n",
    "                elif source_ele == \"bias\":\n",
    "                    if self.struct_param[k][1] not in [\"Simple_Layer\"]:\n",
    "                        item = layer.bias\n",
    "                    else:\n",
    "                        item = layer.b_core\n",
    "                if mode == \"L1\":\n",
    "                    reg = reg + item.abs().sum()\n",
    "                elif mode == \"L2\":\n",
    "                    reg = reg + (item ** 2).sum()\n",
    "                else:\n",
    "                    raise Exception(\"mode {0} not recognized!\".format(mode))\n",
    "        return reg\n",
    "\n",
    "\n",
    "    def get_weights_bias(self, W_source = \"core\", b_source = \"core\"):\n",
    "        W_list = []\n",
    "        b_list = []\n",
    "        for k in range(self.num_layers):\n",
    "            if self.struct_param[k][1] == \"Simple_Layer\":\n",
    "                layer = getattr(self, \"layer_{0}\".format(k))\n",
    "                if W_source == \"core\":\n",
    "                    W_list.append(to_np_array(layer.W_core))\n",
    "                if b_source == \"core\":\n",
    "                    b_list.append(to_np_array(layer.b_core))\n",
    "            elif self.struct_param[k][1] in self.param_available:\n",
    "                layer = getattr(self, \"layer_{0}\".format(k))\n",
    "                if W_source == \"core\":\n",
    "                    W_list.append(to_np_array(layer.weight))\n",
    "                if b_source == \"core\":\n",
    "                    b_list.append(to_np_array(layer.bias, full_reduce = False))\n",
    "            else:\n",
    "                if W_source == \"core\":\n",
    "                    W_list.append(None)\n",
    "                if b_source == \"core\":\n",
    "                    b_list.append(None)\n",
    "        return W_list, b_list\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": self.__class__.__name__}\n",
    "        model_dict[\"net_type\"] = self.__class__.__name__\n",
    "        model_dict[\"input_channels\"] = self.input_channels\n",
    "        model_dict[\"struct_param\"] = self.struct_param\n",
    "        model_dict[\"settings\"] = self.settings\n",
    "        model_dict[\"weights\"], model_dict[\"bias\"] = self.get_weights_bias(W_source = \"core\", b_source = \"core\")\n",
    "        model_dict[\"return_indices\"] = self.return_indices\n",
    "        return model_dict\n",
    "\n",
    "    \n",
    "    @property\n",
    "    def output_size(self):\n",
    "        return self.struct_param[-1][0]\n",
    "    \n",
    "    \n",
    "    @property\n",
    "    def structure(self):\n",
    "        structure = OrderedDict()\n",
    "        structure[\"input_channels\"] = self.input_channels\n",
    "        structure[\"output_size\"] = self.output_size\n",
    "        structure[\"struct_param\"] = self.struct_param if hasattr(self, \"struct_param\") else None\n",
    "        return structure\n",
    "        \n",
    "\n",
    "\n",
    "    def get_sympy_expression(self, verbose=True):\n",
    "        expressions = {i: None for i in range(self.num_layers)}\n",
    "        return expressions\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "    \n",
    "    \n",
    "    def DL(self):\n",
    "        DL = 0\n",
    "        for k in range(self.num_layers):\n",
    "            layer_type = self.struct_param[k][1]\n",
    "            if layer_type in self.param_available:\n",
    "                layer = getattr(self, \"layer_{0}\".format(k))\n",
    "                if layer_type == \"Simple_Layer\":\n",
    "                    DL += layer.DL\n",
    "                else:\n",
    "                    DL += get_list_DL(to_np_array(layer.weight), \"non-snapped\")\n",
    "                    DL += get_list_DL(to_np_array(layer.bias), \"non-snapped\")\n",
    "        return DL\n",
    "    \n",
    "\n",
    "    def load_model_dict(self, model_dict):\n",
    "        new_net = load_model_dict_net(model_dict, is_cuda = self.is_cuda)\n",
    "        self.__dict__.update(new_net.__dict__)\n",
    "\n",
    "\n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        pred_prob = self(X)\n",
    "        if self.return_indices:\n",
    "            pred_prob = pred_prob[0]\n",
    "        pred = pred_prob.max(1)[1]\n",
    "#         self.info_dict[\"accuracy\"] = get_accuracy(pred, y)\n",
    "        return deepcopy(self.info_dict)\n",
    "    \n",
    "\n",
    "    def set_cuda(self, is_cuda):\n",
    "        if isinstance(is_cuda, str):\n",
    "            self.cuda(is_cuda)\n",
    "        else:\n",
    "            if is_cuda:\n",
    "                self.cuda()\n",
    "            else:\n",
    "                self.cpu()\n",
    "        self.is_cuda = is_cuda\n",
    "\n",
    "\n",
    "\n",
    "    def set_trainable(self, is_trainable):\n",
    "        for k in range(self.num_layers):\n",
    "            layer = getattr(self, \"layer_{0}\".format(k))\n",
    "            if self.struct_param[k][1] == \"Simple_Layer\":\n",
    "                layer.set_trainable(is_trainable)\n",
    "            elif self.struct_param[k][1] in self.param_available:\n",
    "                for param in layer.parameters():\n",
    "                    param.requires_grad = is_trainable\n",
    "\n",
    "\n",
    "\n",
    "class Conv_Model(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        encoder_model_dict,\n",
    "        core_model_dict,\n",
    "        decoder_model_dict,\n",
    "        latent_size = 2,\n",
    "        is_generative = True,\n",
    "        is_res_block = True,\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        \"\"\"Conv_Model consists of an encoder, a core and a decoder\"\"\"\n",
    "        super(Conv_Model, self).__init__()\n",
    "        self.latent_size = latent_size\n",
    "        self.is_generative = is_generative\n",
    "        if not is_generative:\n",
    "            self.encoder = load_model_dict(encoder_model_dict, is_cuda = is_cuda)\n",
    "        self.core = load_model_dict(core_model_dict, is_cuda = is_cuda)\n",
    "        self.decoder = load_model_dict(decoder_model_dict, is_cuda = is_cuda)\n",
    "        self.is_res_block = is_res_block\n",
    "        self.is_cuda = is_cuda\n",
    "        self.info_dict = {}\n",
    "\n",
    "\n",
    "    @property\n",
    "    def num_layers(self):\n",
    "        if self.is_generative:\n",
    "            return 1\n",
    "        else:\n",
    "            return len(self.core.model_dict[\"struct_param\"])\n",
    "\n",
    "\n",
    "    def forward(\n",
    "        self,\n",
    "        X,\n",
    "        latent = None,\n",
    "        **kwargs\n",
    "        ):\n",
    "        if self.is_generative:\n",
    "            if len(latent.shape) == 1:\n",
    "                latent = latent.repeat(len(X), 1)\n",
    "            latent = self.core(latent)\n",
    "        else:\n",
    "            p_dict = {k: latent if k == 0 else None for k in range(self.num_layers)}\n",
    "            latent = self.encoder(X)\n",
    "            latent = self.core(latent, p_dict = p_dict)\n",
    "        output = self.decoder(latent)\n",
    "        if self.is_res_block:\n",
    "            output = (X + nn.Sigmoid()(output)).clamp(0, 1)\n",
    "        return output\n",
    "    \n",
    "    \n",
    "    def forward_multistep(self, X, latents, isplot = False, num_images = 1):\n",
    "        assert len(latents.shape) == 1\n",
    "        length = int(len(latents) / 2)\n",
    "        output = X\n",
    "        for i in range(length - 1):\n",
    "            latent = latents[i * self.latent_size: (i + 2) * self.latent_size]\n",
    "            output = self(output, latent = latent)\n",
    "            if isplot:\n",
    "                plot_matrices(output[:num_images,0])\n",
    "        return output\n",
    "\n",
    "\n",
    "    def get_loss(self, X, y, criterion, **kwargs):\n",
    "        return criterion(self(X = X[0], latent = X[1]), y)\n",
    "    \n",
    "    \n",
    "    def plot(self, X, y, num_images = 1):\n",
    "        y_pred = self(X[0], latent = X[1])\n",
    "        idx_list = np.random.choice(len(X[0]), num_images)\n",
    "        for idx in idx_list:\n",
    "            matrix = torch.cat([X[0][idx], y[idx], y_pred[idx]])\n",
    "            plot_matrices(matrix, images_per_row = 8)\n",
    "    \n",
    "    \n",
    "    def get_regularization(self, source = [\"weights\", \"bias\"], mode = \"L1\"):\n",
    "        if self.is_generative:\n",
    "            return self.core.get_regularization(source = source, mode = mode) + \\\n",
    "                    self.decoder.get_regularization(source = source, mode = mode)\n",
    "        else:\n",
    "            return self.encoder.get_regularization(source = source, mode = mode) + \\\n",
    "                    self.core.get_regularization(source = source, mode = mode) + \\\n",
    "                    self.decoder.get_regularization(source = source, mode = mode)\n",
    "\n",
    "\n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return deepcopy(self.info_dict)\n",
    "    \n",
    "    \n",
    "    def set_trainable(self, is_trainable):\n",
    "        if not self.is_generative:\n",
    "            self.encoder.set_trainable(is_trainable)\n",
    "        self.core.set_trainable(is_trainable)\n",
    "        self.decoder.set_trainable(is_trainable)\n",
    "    \n",
    "    \n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Conv_Model\"}\n",
    "        if not self.is_generative:\n",
    "            model_dict[\"encoder_model_dict\"] = self.encoder.model_dict\n",
    "        model_dict[\"latent_size\"] = self.latent_size\n",
    "        model_dict[\"core_model_dict\"] = self.core.model_dict\n",
    "        model_dict[\"decoder_model_dict\"] = self.decoder.model_dict\n",
    "        model_dict[\"is_generative\"] = self.is_generative\n",
    "        model_dict[\"is_res_block\"] = self.is_res_block\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "\n",
    "\n",
    "class Conv_Autoencoder(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_channels_encoder,\n",
    "        input_channels_decoder,\n",
    "        struct_param_encoder,\n",
    "        struct_param_decoder,\n",
    "        latent_size = (1,2),\n",
    "        share_model_among_steps = False,\n",
    "        settings = {},\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        \"\"\"Conv_Autoencoder consists of an encoder and a decoder\"\"\"\n",
    "        super(Conv_Autoencoder, self).__init__()\n",
    "        self.input_channels_encoder = input_channels_encoder\n",
    "        self.input_channels_decoder = input_channels_decoder\n",
    "        self.struct_param_encoder = struct_param_encoder\n",
    "        self.struct_param_decoder = struct_param_decoder\n",
    "        self.share_model_among_steps = share_model_among_steps\n",
    "        self.settings = settings\n",
    "        self.encoder = ConvNet(input_channels = input_channels_encoder, struct_param = struct_param_encoder, settings = settings, is_cuda = is_cuda)\n",
    "        self.decoder = ConvNet(input_channels = input_channels_decoder, struct_param = struct_param_decoder, settings = settings, is_cuda = is_cuda)\n",
    "        self.is_cuda = is_cuda\n",
    "    \n",
    "    def encode(self, input):\n",
    "        if self.share_model_among_steps:\n",
    "            latent = []\n",
    "            for i in range(input.shape[1]):\n",
    "                latent_step = self.encoder(input[:, i:i+1])\n",
    "                latent.append(latent_step)\n",
    "            return torch.cat(latent, 1)\n",
    "        else:\n",
    "            return self.encoder(input)\n",
    "    \n",
    "    def decode(self, latent):\n",
    "        if self.share_model_among_steps:\n",
    "            latent_size = self.struct_param_encoder[-1][0]\n",
    "            latent = latent.view(latent.size(0), -1, latent_size)\n",
    "            output = []\n",
    "            for i in range(latent.shape[1]):\n",
    "                output_step = self.decoder(latent[:, i].contiguous())\n",
    "                output.append(output_step)\n",
    "            return torch.cat(output, 1)\n",
    "        else:\n",
    "            return self.decoder(latent)\n",
    "    \n",
    "    def set_trainable(self, is_trainable):\n",
    "        self.encoder.set_trainable(is_trainable)\n",
    "        self.decoder.set_trainable(is_trainable)\n",
    "    \n",
    "    def forward(self, input):\n",
    "        return self.decode(self.encode(input))\n",
    "    \n",
    "    def get_loss(self, input, target, criterion, **kwargs):\n",
    "        return criterion(self(input), target)\n",
    "    \n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\"):\n",
    "        return self.encoder.get_regularization(source = source, mode = mode) + \\\n",
    "               self.decoder.get_regularization(source = source, mode = mode)\n",
    "    \n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Conv_Autoencoder\"}\n",
    "        model_dict[\"net_type\"] = \"Conv_Autoencoder\"\n",
    "        model_dict[\"input_channels_encoder\"] = self.input_channels_encoder\n",
    "        model_dict[\"input_channels_decoder\"] = self.input_channels_decoder\n",
    "        model_dict[\"struct_param_encoder\"] = self.struct_param_encoder\n",
    "        model_dict[\"struct_param_decoder\"] = self.struct_param_decoder\n",
    "        model_dict[\"share_model_among_steps\"] = self.share_model_among_steps\n",
    "        model_dict[\"settings\"] = self.settings\n",
    "        model_dict[\"encoder\"] = self.encoder.model_dict\n",
    "        model_dict[\"decoder\"] = self.decoder.model_dict\n",
    "        return model_dict\n",
    "    \n",
    "    def load_model_dict(self, model_dict):\n",
    "        model = load_model_dict(model_dict, is_cuda = self.is_cuda)\n",
    "        self.__dict__.update(model.__dict__)\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "    def DL(self):\n",
    "        return self.encoder.DL + self.decoder.DL\n",
    "\n",
    "\n",
    "\n",
    "class Flatten(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Flatten, self).__init__()\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x.view(x.size(0), -1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### VAE:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VAE(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        encoder_model_dict,\n",
    "        decoder_model_dict,\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        super(VAE, self).__init__()\n",
    "        self.encoder = load_model_dict(encoder_model_dict, is_cuda = is_cuda)\n",
    "        self.decoder = load_model_dict(decoder_model_dict, is_cuda = is_cuda)\n",
    "        self.is_cuda = is_cuda\n",
    "        self.info_dict = {}\n",
    "\n",
    "\n",
    "    def encode(self, X):\n",
    "        Z = self.encoder(X)\n",
    "        latent_size = int(Z.shape[-1] / 2)\n",
    "        mu = Z[..., :latent_size]\n",
    "        logvar = Z[..., latent_size:]\n",
    "        return mu, logvar\n",
    "\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = torch.exp(0.5*logvar)\n",
    "        eps = torch.randn_like(std)\n",
    "        return eps.mul(std).add_(mu)\n",
    "\n",
    "\n",
    "    def decode(self, Z):\n",
    "        return self.decoder(Z)\n",
    "\n",
    "\n",
    "    def forward(self, X):\n",
    "        mu, logvar = self.encode(X)\n",
    "        Z = self.reparameterize(mu, logvar)\n",
    "        return self.decode(Z), mu, logvar\n",
    "\n",
    "\n",
    "    def get_loss(self, X, y = None, **kwargs):\n",
    "        recon_X, mu, logvar = self(X)\n",
    "        BCE = F.binary_cross_entropy(recon_X.view(recon_X.shape[0], -1), X.view(X.shape[0], -1), reduction='sum')\n",
    "        # see Appendix B from VAE paper:\n",
    "        # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014\n",
    "        # https://arxiv.org/abs/1312.6114\n",
    "        # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)\n",
    "        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n",
    "        loss = (BCE + KLD) / len(X)\n",
    "        self.info_dict[\"KLD\"] = KLD.item() / len(X)\n",
    "        self.info_dict[\"BCE\"] = BCE.item() / len(X)\n",
    "        return loss\n",
    "\n",
    "\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"VAE\"}\n",
    "        model_dict[\"encoder_model_dict\"] = self.encoder.model_dict\n",
    "        model_dict[\"decoder_model_dict\"] = self.decoder.model_dict\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\"):\n",
    "        return self.encoder.get_regularization(source = source, mode = mode) + self.decoder.get_regularization(source = source, mode = mode)\n",
    "\n",
    "\n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return deepcopy(self.info_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reparameterization toolkit:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net_reparam(nn.Module):\n",
    "    \"\"\"Module that uses reparameterization to take into two inputs and gets a scaler\"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        model_dict,\n",
    "        reparam_mode,\n",
    "        is_cuda=False,\n",
    "        ):\n",
    "        super(Net_reparam, self).__init__()\n",
    "        self.model = load_model_dict(model_dict, is_cuda=is_cuda)\n",
    "        self.reparam_mode = reparam_mode\n",
    "\n",
    "    def forward(self, X, Z, is_outer=False):\n",
    "        \"\"\"\n",
    "        Obtaining single value using reparameterization.\n",
    "\n",
    "        Args:\n",
    "            X shape: [Bx, ...]\n",
    "            Z shape: [S, Bz, Z]\n",
    "            is_outer: whether to use outer product to get a tensor with shape [S, Bz, Bx].\n",
    "        \n",
    "        Returns:\n",
    "            If is_outer==True, return log_prob of shape [S, Bz, Bx]\n",
    "            If is_outer==False, return log_prob of shape [S, Bz]  (where Bz=Bx)\n",
    "        \"\"\"\n",
    "        dist, _ = reparameterize(self.model, X, mode=self.reparam_mode)\n",
    "        if is_outer:\n",
    "            log_prob = dist.log_prob(Z[...,None,:])\n",
    "        else:\n",
    "            log_prob = dist.log_prob(Z)\n",
    "        if self.reparam_mode == 'diag':\n",
    "            log_prob = log_prob.sum(-1)\n",
    "        return log_prob\n",
    "\n",
    "    def get_regularization(self, source = [\"weight\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        return self.model.get_regularization(source=source, model=mode, **kwargs)\n",
    "\n",
    "    def prepare_inspection(self, X, y, **kwargs):\n",
    "        return {}\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Net_reparam\"}\n",
    "        model_dict[\"model\"] = self.model.model_dict\n",
    "        model_dict[\"reparam_mode\"] = self.reparam_mode\n",
    "        return model_dict\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "\n",
    "def reparameterize(model, input, mode=\"full\", size=None):\n",
    "    if mode.startswith(\"diag\"):\n",
    "        if model is not None and model.__class__.__name__ == \"Mixture_Model\":\n",
    "            return reparameterize_mixture_diagonal(model, input, mode=mode)\n",
    "        else:\n",
    "            return reparameterize_diagonal(model, input, mode=mode)\n",
    "    elif mode == \"full\":\n",
    "        return reparameterize_full(model, input, size=size)\n",
    "    else:\n",
    "        raise Exception(\"Mode {} is not valid!\".format(mode))\n",
    "\n",
    "\n",
    "def reparameterize_diagonal(model, input, mode):\n",
    "    if model is not None:\n",
    "        mean_logit = model(input)\n",
    "    else:\n",
    "        mean_logit = input\n",
    "    if mode.startswith(\"diagg\"):\n",
    "        if isinstance(mean_logit, tuple):\n",
    "            mean = mean_logit[0]\n",
    "        else:\n",
    "            mean = mean_logit\n",
    "        std = torch.ones(mean.shape).to(mean.device)\n",
    "        dist = Normal(mean, std)\n",
    "        return dist, (mean, std)\n",
    "    elif mode.startswith(\"diag\"):\n",
    "        if isinstance(mean_logit, tuple):\n",
    "            mean_logit = mean_logit[0]\n",
    "        size = int(mean_logit.size(-1) / 2)\n",
    "        mean = mean_logit[:, :size]\n",
    "        std = F.softplus(mean_logit[:, size:], beta=1) + 1e-10\n",
    "        dist = Normal(mean, std)\n",
    "        return dist, (mean, std)\n",
    "    else:\n",
    "        raise Exception(\"mode {} is not valid!\".format(mode))\n",
    "\n",
    "\n",
    "def reparameterize_mixture_diagonal(model, input, mode):\n",
    "    mean_logit, weight_logits = model(input)\n",
    "    if mode.startswith(\"diagg\"):\n",
    "        mean_list = mean_logit\n",
    "        scale_list = torch.ones(mean_list.shape).to(mean_list.device)\n",
    "    else:\n",
    "        size = int(mean_logit.size(-2) / 2)\n",
    "        mean_list = mean_logit[:, :size]\n",
    "        scale_list = F.softplus(mean_logit[:, size:], beta=1) + 0.01  # Avoid the std to go to 0\n",
    "    dist = Mixture_Gaussian_reparam(mean_list=mean_list,\n",
    "                                    scale_list=scale_list,\n",
    "                                    weight_logits=weight_logits,\n",
    "                                   )\n",
    "    return dist, (mean_list, scale_list)\n",
    "\n",
    "\n",
    "def reparameterize_full(model, input, size=None):\n",
    "    if model is not None:\n",
    "        mean_logit = model(input)\n",
    "    else:\n",
    "        mean_logit = input\n",
    "    if isinstance(mean_logit, tuple):\n",
    "        mean_logit = mean_logit[0]\n",
    "    if size is None:\n",
    "        dim = mean_logit.size(-1)\n",
    "        size = int((np.sqrt(9 + 8 * dim) - 3) / 2)\n",
    "    mean = mean_logit[:, :size]\n",
    "    scale_tril = fill_triangular(mean_logit[:, size:], size)\n",
    "    scale_tril = matrix_diag_transform(scale_tril, F.softplus)\n",
    "    dist = MultivariateNormal(mean, scale_tril = scale_tril)\n",
    "    return dist, (mean, scale_tril)\n",
    "\n",
    "\n",
    "def sample(dist, n=None):\n",
    "    \"\"\"Sample n instances from distribution dist\"\"\"\n",
    "    if n is None:\n",
    "        return dist.rsample()\n",
    "    else:\n",
    "        return dist.rsample((n,))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Probability models:\n",
    "### Mixture of Gaussian:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mixture_Gaussian(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        num_components,\n",
    "        dim,\n",
    "        param_mode = \"full\",\n",
    "        is_cuda = False,\n",
    "        ):\n",
    "        super(Mixture_Gaussian, self).__init__()\n",
    "        self.num_components = num_components\n",
    "        self.dim = dim\n",
    "        self.param_mode = param_mode\n",
    "        self.is_cuda = is_cuda\n",
    "        self.device = torch.device(self.is_cuda if isinstance(self.is_cuda, str) else \"cuda\" if self.is_cuda else \"cpu\")\n",
    "        self.info_dict = {}\n",
    "\n",
    "\n",
    "    def initialize(self, model_dict = None, input = None, num_samples = 100, verbose = False):\n",
    "        if input is not None:\n",
    "            neg_log_prob_min = np.inf\n",
    "            loc_init_min = None\n",
    "            scale_init_min = None\n",
    "            for i in range(num_samples):\n",
    "                neg_log_prob, loc_init_list, scale_init_list = self.initialize_ele(input)\n",
    "                if verbose:\n",
    "                    print(\"{0}: neg_log_prob: {1:.4f}\".format(i, neg_log_prob))\n",
    "                if neg_log_prob < neg_log_prob_min:\n",
    "                    neg_log_prob_min = neg_log_prob\n",
    "                    loc_init_min = self.loc_list.detach()\n",
    "                    scale_init_min = self.scale_list.detach()\n",
    "\n",
    "            self.loc_list = nn.Parameter(loc_init_min.to(self.device))\n",
    "            self.scale_list = nn.Parameter(scale_init_min.to(self.device))\n",
    "            print(\"min neg_log_prob: {0:.6f}\".format(to_np_array(neg_log_prob_min)))\n",
    "        else:\n",
    "            if model_dict is None:\n",
    "                self.weight_logits = nn.Parameter((torch.randn(self.num_components) * np.sqrt(2 / (1 + self.dim))).to(self.device))\n",
    "            else:\n",
    "                self.weight_logits = nn.Parameter((torch.FloatTensor(model_dict[\"weight_logits\"])).to(self.device))\n",
    "            if self.param_mode == \"full\": \n",
    "                size = self.dim * (self.dim + 1) // 2\n",
    "            elif self.param_mode == \"diag\":\n",
    "                size = self.dim\n",
    "            else:\n",
    "                raise\n",
    "            \n",
    "            if model_dict is None:\n",
    "                self.loc_list = nn.Parameter(torch.randn(self.num_components, self.dim).to(self.device))\n",
    "                self.scale_list = nn.Parameter((torch.randn(self.num_components, size) / self.dim).to(self.device))\n",
    "            else:\n",
    "                self.loc_list = nn.Parameter(torch.FloatTensor(model_dict[\"loc_list\"]).to(self.device))\n",
    "                self.scale_list = nn.Parameter(torch.FloatTensor(model_dict[\"scale_list\"]).to(self.device))\n",
    "\n",
    "\n",
    "    def initialize_ele(self, input):\n",
    "        if self.param_mode == \"full\":\n",
    "            size = self.dim * (self.dim + 1) // 2\n",
    "        elif self.param_mode == \"diag\":\n",
    "            size = self.dim\n",
    "        else:\n",
    "            raise\n",
    "        length = len(input)\n",
    "        self.weight_logits = nn.Parameter(torch.zeros(self.num_components).to(self.device))\n",
    "        self.loc_list = nn.Parameter(input[torch.multinomial(torch.ones(length) / length, self.num_components)].detach())\n",
    "        self.scale_list = nn.Parameter((torch.randn(self.num_components, size).to(self.device) * input.std() / 5).to(self.device))\n",
    "        neg_log_prob = self.get_loss(input)\n",
    "        return neg_log_prob\n",
    "\n",
    "\n",
    "    def prob(self, input):\n",
    "        if len(input.shape) == 1:\n",
    "            input = input.unsqueeze(1)\n",
    "        assert len(input.shape) in [0, 2, 3]\n",
    "        input = input.unsqueeze(-2)\n",
    "        if self.param_mode == \"diag\":\n",
    "            scale_list = F.softplus(self.scale_list)\n",
    "            logits = (- (input - self.loc_list) ** 2 / 2 / scale_list ** 2 - torch.log(scale_list * np.sqrt(2 * np.pi))).sum(-1)\n",
    "        else:\n",
    "            raise\n",
    "        prob = torch.matmul(torch.exp(logits), nn.Softmax(dim = 0)(self.weight_logits))\n",
    "#         prob_list = []\n",
    "#         for i in range(self.num_components):\n",
    "#             if self.param_mode == \"full\":\n",
    "#                 scale_tril = fill_triangular(getattr(self, \"scale_{0}\".format(i)), self.dim)\n",
    "#                 scale_tril = matrix_diag_transform(scale_tril, F.softplus)\n",
    "#                 dist = MultivariateNormal(getattr(self, \"loc_{0}\".format(i)), scale_tril = scale_tril)\n",
    "#                 log_prob = dist.log_prob(input)\n",
    "#             elif self.param_mode == \"diag\":\n",
    "#                 dist = Normal(getattr(self, \"loc_{0}\".format(i)).unsqueeze(0), F.softplus(getattr(self, \"scale_{0}\".format(i))))\n",
    "#                 mu = getattr(self, \"loc_{0}\".format(i)).unsqueeze(0)\n",
    "#                 sigma = F.softplus(getattr(self, \"scale_{0}\".format(i)))\n",
    "#                 log_prob = (- (input - mu) ** 2 / 2 / sigma ** 2 - torch.log(sigma * np.sqrt(2 * np.pi))).sum(-1)\n",
    "#             else:\n",
    "#                 raise\n",
    "#             setattr(self, \"component_{0}\".format(i), dist)\n",
    "#             prob = torch.exp(log_prob)\n",
    "#             prob_list.append(prob)\n",
    "#         prob_list = torch.stack(prob_list, -1)\n",
    "#         prob = torch.matmul(prob_list, nn.Softmax(dim = 0)(self.weight_logits))\n",
    "        return prob\n",
    "\n",
    "\n",
    "    def log_prob(self, input):\n",
    "        return torch.log(self.prob(input) + 1e-45)\n",
    "\n",
    "\n",
    "    def get_loss(self, X, y = None, **kwargs):\n",
    "        \"\"\"Optimize negative log-likelihood\"\"\"\n",
    "        neg_log_prob = - self.log_prob(X).mean() / np.log(2)\n",
    "        self.info_dict[\"loss\"] = to_np_array(neg_log_prob)\n",
    "        return neg_log_prob\n",
    "\n",
    "\n",
    "    def prepare_inspection(X, y, criterion, **kwargs):\n",
    "        return deepcopy(self.info_dict)\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Mixture_Gaussian\"}\n",
    "        model_dict[\"num_components\"] = self.num_components\n",
    "        model_dict[\"dim\"] = self.dim\n",
    "        model_dict[\"param_mode\"] = self.param_mode\n",
    "        model_dict[\"weight_logits\"] = to_np_array(self.weight_logits)\n",
    "        model_dict[\"loc_list\"] = to_np_array(self.loc_list)\n",
    "        model_dict[\"scale_list\"] = to_np_array(self.scale_list)\n",
    "        return model_dict\n",
    "\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)\n",
    "\n",
    "\n",
    "    def get_param(self):\n",
    "        weights = to_np_array(nn.Softmax(dim = 0)(self.weight_logits))\n",
    "        loc_list = to_np_array(self.loc_list)\n",
    "        scale_list = to_np_array(self.scale_list)\n",
    "        print(\"weights: {0}\".format(weights))\n",
    "        print(\"loc:\")\n",
    "        pp.pprint(loc_list)\n",
    "        print(\"scale:\")\n",
    "        pp.pprint(scale_list)\n",
    "        return weights, loc_list, scale_list\n",
    "\n",
    "\n",
    "    def visualize(self, input):\n",
    "        import scipy\n",
    "        import matplotlib.pylab as plt\n",
    "        std = to_np_array(input.std())\n",
    "        X = np.arange(to_np_array(input.min()) - 0.2 * std, to_np_array(input.max()) + 0.2 * std, 0.1)\n",
    "        Y_dict = {}\n",
    "        weights = nn.Softmax(dim = 0)(self.weight_logits)\n",
    "        plt.figure(figsize=(10, 4), dpi=100).set_facecolor('white')\n",
    "        for i in range(self.num_components):\n",
    "            Y_dict[i] = weights[0].item() * scipy.stats.norm.pdf((X - self.loc_list[i].item()) / self.scale_list[i].item())\n",
    "            plt.plot(X, Y_dict[i])\n",
    "        Y = np.sum([item for item in Y_dict.values()], 0)\n",
    "        plt.plot(X, Y, 'k--')\n",
    "        plt.plot(input.data.numpy(), np.zeros(len(input)), 'k*')\n",
    "        plt.title('Density of {0}-component mixture model'.format(self.num_components))\n",
    "        plt.ylabel('probability density');\n",
    "\n",
    "\n",
    "    def get_regularization(self, source = [\"weights\", \"bias\"], mode = \"L1\", **kwargs):\n",
    "        reg = to_Variable([0], requires_grad = False).to(self.device)\n",
    "        return reg"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mixture_Gaussian for reparameterization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mixture_Gaussian_reparam(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        # Use as reparamerization:\n",
    "        mean_list=None,\n",
    "        scale_list=None,\n",
    "        weight_logits=None,\n",
    "        # Use as prior:\n",
    "        Z_size=None,\n",
    "        n_components=None,\n",
    "        mean_scale=0.1,\n",
    "        scale_scale=0.1,\n",
    "        # Mode:\n",
    "        is_reparam=True,\n",
    "        reparam_mode=\"diag\",\n",
    "        is_cuda=False,\n",
    "    ):\n",
    "        super(Mixture_Gaussian_reparam, self).__init__()\n",
    "        self.is_reparam = is_reparam\n",
    "        self.reparam_mode = reparam_mode\n",
    "        self.is_cuda = is_cuda\n",
    "        self.device = torch.device(self.is_cuda if isinstance(self.is_cuda, str) else \"cuda\" if self.is_cuda else \"cpu\")\n",
    "\n",
    "        if self.is_reparam:\n",
    "            self.mean_list = mean_list         # size: [B, Z, k]\n",
    "            self.scale_list = scale_list       # size: [B, Z, k]\n",
    "            self.weight_logits = weight_logits # size: [B, k]\n",
    "            self.n_components = self.weight_logits.shape[-1]\n",
    "            self.Z_size = self.mean_list.shape[-2]\n",
    "        else:\n",
    "            self.n_components = n_components\n",
    "            self.Z_size = Z_size\n",
    "            self.mean_list = nn.Parameter((torch.rand(1, Z_size, n_components) - 0.5) * mean_scale)\n",
    "            if reparam_mode == \"diag\":\n",
    "                self.scale_list = nn.Parameter(torch.log(torch.exp((torch.rand(1, Z_size, n_components) * 0.2 + 0.9) * scale_scale) - 1))\n",
    "            elif reparam_mode == \"diagg\":\n",
    "                self.register_buffer('scale_list', torch.ones(1, Z_size, n_components))\n",
    "            else:\n",
    "                raise\n",
    "            self.weight_logits = nn.Parameter(torch.zeros(1, n_components))\n",
    "            if mean_list is not None:\n",
    "                self.mean_list.data = to_Variable(mean_list)\n",
    "                self.scale_list.data = to_Variable(scale_list)\n",
    "                self.weight_logits.data = to_Variable(weight_logits)\n",
    "\n",
    "        self.to(self.device)\n",
    "\n",
    "\n",
    "    def log_prob(self, input):\n",
    "        \"\"\"Obtain the log_prob of the input.\"\"\"\n",
    "        input = input.unsqueeze(-1)  # [S, B, Z, 1]\n",
    "        if self.reparam_mode in [\"diag\", \"diagg\"]:\n",
    "            if self.is_reparam:\n",
    "                # logits: [S, B, Z, k]\n",
    "                logits = - (input - self.mean_list) ** 2 / 2 / self.scale_list ** 2 - torch.log(self.scale_list * np.sqrt(2 * np.pi))\n",
    "            else:\n",
    "                scale_list = F.softplus(self.scale_list, beta=1)\n",
    "                logits = - (input - self.mean_list) ** 2 / 2 / scale_list ** 2 - torch.log(scale_list * np.sqrt(2 * np.pi))\n",
    "        else:\n",
    "            raise\n",
    "        # log_softmax(weight_logits): [B, k]\n",
    "        # logits: [S, B, Z, k]\n",
    "        # log_prob: [S, B, Z]\n",
    "        log_prob = torch.logsumexp(logits + F.log_softmax(self.weight_logits, -1).unsqueeze(-2), axis=-1)  # F(...).unsqueeze(-2): [B, 1, k]\n",
    "        return log_prob\n",
    "\n",
    "\n",
    "    def prob(self, Z):\n",
    "        return torch.exp(self.log_prob(Z))\n",
    "\n",
    "\n",
    "    def sample(self, n=None):\n",
    "        if n is None:\n",
    "            n_core = 1\n",
    "        else:\n",
    "            assert isinstance(n, tuple)\n",
    "            n_core = n[0]\n",
    "        weight_probs = F.softmax(self.weight_logits, -1)  # size: [B, m]\n",
    "        idx = torch.multinomial(weight_probs, n_core, replacement=True).unsqueeze(-2).expand(-1, self.mean_list.shape[-2], -1)  # multinomial result: [B, S]; result: [B, Z, S]\n",
    "        mean_list  = torch.gather(self.mean_list,  dim=-1, index=idx)  # [B, Z, S]\n",
    "        if self.is_reparam:\n",
    "            scale_list = torch.gather(self.scale_list, dim=-1, index=idx)  # [B, Z, S]\n",
    "        else:\n",
    "            scale_list = F.softplus(torch.gather(self.scale_list, dim=-1, index=idx), beta=1)  # [B, Z, S]\n",
    "        Z = torch.normal(mean_list, scale_list).permute(2, 0, 1)\n",
    "        if n is None:\n",
    "            Z = Z.squeeze(0)\n",
    "        return Z\n",
    "\n",
    "\n",
    "    def rsample(self, n=None):\n",
    "        return self.sample(n=n)\n",
    "\n",
    "\n",
    "    def __repr__(self):\n",
    "        return \"Mixture_Gaussian_reparam({}, Z_size={})\".format(self.n_components, self.Z_size)\n",
    "\n",
    "\n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Mixture_Gaussian_reparam\"}\n",
    "        model_dict[\"is_reparam\"] = self.is_reparam\n",
    "        model_dict[\"reparam_mode\"] = self.reparam_mode\n",
    "        model_dict[\"Z_size\"] = self.Z_size\n",
    "        model_dict[\"n_components\"] = self.n_components\n",
    "        model_dict[\"mean_list\"] = to_np_array(self.mean_list)\n",
    "        model_dict[\"scale_list\"] = to_np_array(self.scale_list)\n",
    "        model_dict[\"weight_logits\"] = to_np_array(self.weight_logits)\n",
    "        return model_dict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Triangular distribution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Triangular_dist(Distribution):\n",
    "    \"\"\"Probability distribution with a Triangular shape.\"\"\"\n",
    "    def __init__(self, loc, a, b, validate_args=None):\n",
    "        self.loc, self.a, self.b = broadcast_all(loc, a, b)\n",
    "        batch_shape = torch.Size() if isinstance(loc, Number) else self.loc.size()\n",
    "        super(Triangular_dist, self).__init__(batch_shape, validate_args=validate_args)\n",
    "        \n",
    "    @property\n",
    "    def mean(self):\n",
    "        return self.loc + (self.b - self.a) / 3\n",
    "    \n",
    "    @property\n",
    "    def variance(self):\n",
    "        return (self.a ** 2 + self.b ** 2 + self.a * self.b) / 18\n",
    "    \n",
    "    @property\n",
    "    def stddev(self):\n",
    "        return torch.sqrt(self.variance)\n",
    "    \n",
    "    def expand(self, batch_shape, _instance=None):\n",
    "        new = self._get_checked_instance(PieceWise, _instance)\n",
    "        batch_shape = torch.Size(batch_shape)\n",
    "        new.loc = self.loc.expand(batch_shape)\n",
    "        new.a = self.a.expand(batch_shape)\n",
    "        new.b = self.b.expand(batch_shape)\n",
    "        super(Triangular_dist, new).__init__(batch_shape, validate_args=False)\n",
    "        new._validate_args = self._validate_args\n",
    "        return new\n",
    "    \n",
    "    @constraints.dependent_property\n",
    "    def support(self):\n",
    "        return constraints.interval(self.loc - self.a, self.loc + self.b)\n",
    "    \n",
    "    def sample(self, sample_shape=torch.Size()):\n",
    "        shape = self._extended_shape(sample_shape)\n",
    "        rand = torch.rand(shape, dtype=self.loc.dtype, device=self.loc.device)\n",
    "        with torch.no_grad():\n",
    "            return self.icdf(rand)\n",
    "    \n",
    "    def rsample(self, sample_shape=torch.Size()):\n",
    "        \"\"\"Sample with reparameterization.\"\"\"\n",
    "        shape = self._extended_shape(sample_shape)\n",
    "        rand = torch.rand(shape, dtype=self.loc.dtype, device=self.loc.device)\n",
    "        return self.icdf(rand)\n",
    "\n",
    "    def icdf(self, value):\n",
    "        \"\"\"Inverse cdf.\"\"\"\n",
    "        if self._validate_args:\n",
    "            self._validate_sample(value)\n",
    "        assert value.min() >= 0 and value.max() <= 1\n",
    "        value, loc, a, b = broadcast_all(value, self.loc, self.a, self.b)\n",
    "        a_plus_b = a + b\n",
    "        idx = value < a / a_plus_b\n",
    "        iidx = ~idx\n",
    "        out = torch.ones_like(value)\n",
    "        out[idx] = loc[idx] - a[idx] + torch.sqrt(a[idx] * a_plus_b[idx] * value[idx])\n",
    "        out[iidx] = loc[iidx] + b[iidx] - torch.sqrt(b[iidx] * a_plus_b[iidx] * (1 - value[iidx]) )\n",
    "        return out\n",
    "\n",
    "    def prob(self, value):\n",
    "        \"\"\"Get probability.\"\"\"\n",
    "        if self._validate_args:\n",
    "            self._validate_sample(value)\n",
    "        # compute the variance\n",
    "        value, loc, a, b = broadcast_all(value, self.loc, self.a, self.b)\n",
    "        idx1 = (loc - a <= value) & (value <= loc)\n",
    "        idx2 = (loc < value) & (value <= loc + b)\n",
    "        a_plus_b = a + b\n",
    "\n",
    "        out = torch.zeros_like(value)\n",
    "        out[idx1] = 2 * (value[idx1] - loc[idx1] + a[idx1]) / a[idx1] / a_plus_b[idx1]\n",
    "        out[idx2] = -2 * (value[idx2] - loc[idx2] - b[idx2]) / b[idx2] / a_plus_b[idx2]\n",
    "        return out\n",
    "\n",
    "    def log_prob(self, value):\n",
    "        \"\"\"Get log probability.\"\"\"\n",
    "        return torch.log(self.prob(value))\n",
    "    \n",
    "    @property\n",
    "    def model_dict(self):\n",
    "        model_dict = {\"type\": \"Triangular_dist\"}\n",
    "        model_dict[\"loc\"] = to_np_array(self.loc)\n",
    "        model_dict[\"a\"] = to_np_array(self.a)\n",
    "        model_dict[\"b\"] = to_np_array(self.b)\n",
    "        return model_dict\n",
    "\n",
    "    def load(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        model_dict = load_model(filename, mode=mode)\n",
    "        self.load_model_dict(model_dict)\n",
    "\n",
    "    def save(self, filename):\n",
    "        mode = \"json\" if filename.endswith(\".json\") else \"pickle\"\n",
    "        save_model(self.model_dict, filename, mode=mode)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model_dict_distribution(model_dict, is_cuda = False):\n",
    "    if model_dict[\"type\"] == \"Mixture_Gaussian\":\n",
    "        model = Mixture_Gaussian(\n",
    "            num_components=model_dict[\"num_components\"],\n",
    "            dim=model_dict[\"dim\"],\n",
    "            param_mode=model_dict[\"param_mode\"],\n",
    "            is_cuda=is_cuda,\n",
    "        )\n",
    "        model.initialize(model_dict = model_dict)\n",
    "    elif model_dict[\"type\"] == \"Mixture_Gaussian_reparam\":\n",
    "        model = Mixture_Gaussian_reparam(\n",
    "            is_reparam=model_dict[\"is_reparam\"],\n",
    "            reparam_mode=model_dict[\"reparam_mode\"],\n",
    "            mean_list=model_dict[\"mean_list\"],\n",
    "            scale_list=model_dict[\"scale_list\"],\n",
    "            weight_logits=model_dict[\"weight_logits\"],\n",
    "            Z_size=model_dict[\"Z_size\"],\n",
    "            n_components=model_dict[\"n_components\"],\n",
    "            is_cuda=is_cuda,\n",
    "        )\n",
    "    elif model_dict[\"type\"] == \"Triangular_dist\":\n",
    "        model = Triangular_dist(\n",
    "            loc=model_dict[\"loc\"],\n",
    "            a=model_dict[\"a\"],\n",
    "            b=model_dict[\"b\"],\n",
    "        )\n",
    "    else:\n",
    "        raise Exception(\"Type {} is not valid!\".format(model_dict[\"type\"]))\n",
    "    return model"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.5"
  },
  "toc-autonumbering": false,
  "toc-showcode": false
 },
 "nbformat": 4,
 "nbformat_minor": 4
}