inference_zero_shot.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aadbdac5-f4d8-4cc0-8154-238e64bf7b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "from collections import OrderedDict, Iterable\n",
    "from copy import deepcopy\n",
    "from datetime import datetime\n",
    "import matplotlib.pylab as plt\n",
    "import gc\n",
    "import logging\n",
    "logging.getLogger('matplotlib.font_manager').disabled = True\n",
    "import matplotlib.pylab as plt\n",
    "import os\n",
    "import pdb\n",
    "import pickle\n",
    "import pprint as pp\n",
    "import random\n",
    "import pandas as pd\n",
    "from scipy import stats\n",
    "from IPython.display import display, HTML\n",
    "pd.options.display.max_rows = 1500\n",
    "pd.options.display.max_columns = 200\n",
    "pd.options.display.width = 1000\n",
    "pd.set_option('max_colwidth', 400)\n",
    "\n",
    "from IPython.display import display\n",
    "import numpy as np\n",
    "from sklearn import linear_model\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import datasets, transforms, utils\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..'))\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from zeroc.datasets.arc_image import ARCDataset\n",
    "from zeroc.datasets.BabyARC.code.dataset.dataset import *\n",
    "from zeroc.concept_library.concepts import Concept, Concept_Pattern, Placeholder, Tensor\n",
    "from zeroc.argparser import update_default_hyperparam, get_args_EBM, get_SGLD_kwargs\n",
    "from zeroc.concept_library.models import load_best_model, GraphEBM, neg_mask_sgd, neg_mask_sgd_ensemble, ResBlock, CResBlock, spectral_norm\n",
    "from zeroc.train import get_c_core, init_concepts_with_repr, ConceptDataset, ConceptFewshotDataset, ConceptCompositionDataset, get_dataset, requires_grad, test_acc, load_model_energy, SampleBuffer, sample_buffer, id_to_tensor\n",
    "from zeroc.concept_library.settings import REPR_DIM\n",
    "from zeroc.utils import REA_PATH, EXP_PATH\n",
    "try:\n",
    "    from zeroc.concept_library.util import try_call, plot_simple, plot_2_axis, Attr_Dict, MineDataset, Batch, extend_dims, groupby_add_keys, get_unique_keys_df, filter_df, groupby_add_keys, to_cpu_recur, Printer, get_num_params, transform_dict, get_graph_edit_distance, draw_nx_graph, get_nx_graph, get_triu_ids, get_soft_IoU, get_pdict, Batch, pdump, pload, filter_kwargs, gather_broadcast, ddeepcopy as deepcopy, to_Variable, set_seed, Zip, COLOR_LIST, init_args, make_dir, str2bool, get_filename_short, get_machine_name, get_device, record_data, plot_matrices, filter_filename, get_next_available_key, to_np_array, print_banner, get_filename_short, write_to_config, to_cpu\n",
    "    from zeroc.concept_library.util import find_connected_components_colordiff, onehot_to_RGB, classify_concept, Shared_Param_Dict, get_inputs_targets_EBM, repeat_n, color_dict, to_one_hot, onehot_to_RGB, get_root_dir, get_module_parameters, assign_embedding_value, get_hashing, to_device_recur, visualize_matrices\n",
    "except Exception as e:\n",
    "    raise Exception(f\"{e}. Please update the concept_library by running 'git submodule init; git submodule update' in the local zeroc repo!\")\n",
    "\n",
    "p = Printer()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "104797e9-836f-460d-90cb-e20e0d6bf40d",
   "metadata": {},
   "source": [
    "# 1. Helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ba78830-072f-430d-a8c4-8701da2caf3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_loss(data_record, interval=1):\n",
    "    fontsize = 12\n",
    "    plt.figure(figsize=(8,6))\n",
    "    plt.plot(data_record['epoch'][::interval], data_record['pos_out'][::interval], label=\"pos\")\n",
    "    plt.plot(data_record['epoch'][::interval], data_record['neg_out'][::interval], label=\"neg\")\n",
    "    plt.plot(data_record['epoch'][::interval], data_record['loss'][::interval], label=\"loss\")\n",
    "    plt.plot(data_record['epoch'][::interval], np.array(data_record['loss'][::interval]) - np.array(data_record['pos_out'][::interval]), label=\"loss-pos\")\n",
    "    plt.legend(fontsize=fontsize)\n",
    "    plt.tick_params(labelsize=fontsize)\n",
    "    plt.xlabel(\"epoch\", fontsize=fontsize)\n",
    "    plt.ylabel(\"loss\", fontsize=fontsize)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def update_CONCEPTS_OPERATORS(CONCEPTS, OPERATORS, concept_embeddings_load, update_keys=None):\n",
    "    if update_keys is not None and not isinstance(update_keys, list):\n",
    "        update_keys = [update_keys]\n",
    "    for key in CONCEPTS:\n",
    "        if key in concept_embeddings_load:\n",
    "            if update_keys is not None and key in update_keys and key in concept_embeddings_load:\n",
    "                c_repr_curr = CONCEPTS[key].get_node_repr()\n",
    "                c_repr_curr.data = torch.FloatTensor(concept_embeddings_load[key])\n",
    "    for key in OPERATORS:\n",
    "        if update_keys is not None and key in update_keys and key in concept_embeddings_load:\n",
    "            c_repr_curr = OPERATORS[key].get_node_repr()\n",
    "            c_repr_curr.data = torch.FloatTensor(concept_embeddings_load[key])\n",
    "\n",
    "\n",
    "def test_concept_embedding(CONCEPTS, OPERATORS, concept_embeddings_load, raise_warnings_only=False, checked_keys=None):\n",
    "    concept_embeddings = {key: to_np_array(CONCEPTS[key].get_node_repr()) for key in CONCEPTS}\n",
    "    operator_embeddings = {key: to_np_array(OPERATORS[key].get_node_repr()) for key in OPERATORS}\n",
    "    concept_embeddings.update(operator_embeddings)\n",
    "    different_list = []\n",
    "    same_list = []\n",
    "    for key, value in concept_embeddings_load.items():\n",
    "        if key not in concept_embeddings:\n",
    "            p.warning(\"key '{}' not in the current concept/relation embedding.\".format(key))\n",
    "            continue\n",
    "        if (checked_keys is None or key in checked_keys) and np.abs(value - concept_embeddings[key]).max() > 0:\n",
    "            different_list.append(key)\n",
    "        else:\n",
    "            same_list.append(key)\n",
    "    if len(different_list) > 0:\n",
    "        if raise_warnings_only:\n",
    "            print(\"The c_repr for {} are the same.\\nThe c_repr for {} are different.\\n\".format(same_list, different_list))\n",
    "        else:\n",
    "            raise Exception(\"The c_repr for {} are the same.\\nThe c_repr for {} are different.\\n\".format(same_list, different_list))\n",
    "\n",
    "\n",
    "def get_model_samples(\n",
    "    model,\n",
    "    args,\n",
    "    dataset=None,\n",
    "    init=\"gaussian\",\n",
    "    sample_step=None,\n",
    "    batch_size=None,\n",
    "    ensemble_size=1,\n",
    "    plot_ensemble_mode=\"min\",\n",
    "    analysis_modes=[\"mask|c_repr\", \"c_repr|mask\", \"mask|c\", \"c_repr|c\"],\n",
    "    w_type=\"image+mask\",\n",
    "    plot_grey_scale=False,\n",
    "    concept_collection=None,\n",
    "    CONCEPTS=None,\n",
    "    OPERATORS=None,\n",
    "    isplot=True,\n",
    "    device=\"cpu\",\n",
    "    plot_topk=3,\n",
    "):\n",
    "    \"\"\"This one is with concept_collection argument.\"\"\"\n",
    "    def init_neg_mask(pos_img, buffer, args):\n",
    "        \"\"\"Initialize negative mask\"\"\"\n",
    "        neg_data = sample_buffer(\n",
    "            buffer,\n",
    "            in_channels=args.in_channels,\n",
    "            n_classes=args.n_classes,\n",
    "            image_size=args.image_size,\n",
    "            batch_size=batch_size*ensemble_size,\n",
    "            is_mask=args.is_mask,\n",
    "            is_two_branch=args.is_two_branch,\n",
    "            w_type=args.w_type,\n",
    "            p=args.p_buffer,\n",
    "            device=device,\n",
    "        )\n",
    "        _, neg_mask, _, _ = neg_data\n",
    "        pos_img_l = repeat_n(pos_img, n_repeats=ensemble_size)\n",
    "        if args.is_two_branch:\n",
    "            if init == \"mask-gaussian\":\n",
    "                neg_mask = (neg_mask[0] * (pos_img_l[0].argmax(1)[:, None] != 0), neg_mask[1] * (pos_img_l[1].argmax(1)[:, None] != 0))\n",
    "        else:\n",
    "            if init == \"mask-gaussian\":\n",
    "                pos_img_l = repeat_n(pos_img, n_repeats=ensemble_size)\n",
    "                neg_mask = (neg_mask[0] * (pos_img_l.argmax(1)[:, None] != 0),)\n",
    "        for k in range(len(neg_mask)):\n",
    "            neg_mask[k].requires_grad = True\n",
    "        return neg_mask\n",
    "\n",
    "    model.eval()\n",
    "    buffer = SampleBuffer()\n",
    "    args = deepcopy(args)\n",
    "    batch_size = args.batch_size if batch_size is None else batch_size\n",
    "    sample_step = args.sample_step if sample_step is None else sample_step\n",
    "    if concept_collection is not None:\n",
    "        args.concept_collection = concept_collection\n",
    "    args.sample_step = sample_step\n",
    "    args.ebm_target = \"mask\"\n",
    "\n",
    "    if args.is_mask:\n",
    "        assert dataset is not None\n",
    "        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, collate_fn=Batch(is_collate_tuple=True).collate())\n",
    "        for pos_data in dataloader:\n",
    "            break\n",
    "        pos_img, pos_mask, pos_id, _ = pos_data\n",
    "        pos_repr = id_to_tensor(pos_id, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device)\n",
    "        if args.is_two_branch:\n",
    "            pos_img = (pos_img[0].to(device), pos_img[1].to(device))\n",
    "        else:\n",
    "            pos_img = pos_img.to(device)\n",
    "        pos_mask = to_device_recur(pos_mask, device)\n",
    "\n",
    "        data = {\n",
    "            \"pos_img\": pos_img,\n",
    "            \"pos_mask\": pos_mask,\n",
    "            \"pos_id\": pos_id,\n",
    "            \"pos_repr\": pos_repr,\n",
    "            \"concept_collection\": args.concept_collection,\n",
    "        }\n",
    "\n",
    "        # Calculate E(ground truth mask for given concept)\n",
    "        if \"mask|c_repr\" in analysis_modes:\n",
    "            neg_mask = init_neg_mask(pos_img, buffer, args)\n",
    "            (_, neg_mask_ensemble, _, _, _), neg_out_list_ensemble, _ = neg_mask_sgd_ensemble(\n",
    "                model, pos_img, neg_mask, pos_repr, z=None, zgnn=None, wtarget=None, args=args,\n",
    "                ensemble_size=ensemble_size,\n",
    "            )\n",
    "            data[\"mask|c_repr\"] = {\n",
    "                \"neg_mask_ensemble\": neg_mask_ensemble,\n",
    "                \"neg_out_list_ensemble\": neg_out_list_ensemble,\n",
    "            }\n",
    "        # Calculate E(concept given a ground truth mask)\n",
    "        if \"c_repr|mask\" in analysis_modes:\n",
    "            data[\"c_repr|mask\"] = {\"c_repr_pred\": model.classify(pos_img, pos_mask, args.concept_collection, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS)}\n",
    "\n",
    "        if \"mask|c\" in analysis_modes:\n",
    "            neg_mask_ensemble_collection = []\n",
    "            neg_out_list_ensemble_collection = []\n",
    "            for j in range(len(args.concept_collection)):\n",
    "                neg_mask = init_neg_mask(pos_img, buffer, args)\n",
    "                c_repr = id_to_tensor([args.concept_collection[j]] * len(pos_repr), CONCEPTS=CONCEPTS, OPERATORS=OPERATORS).to(device)\n",
    "                (_, neg_mask_ensemble, _, _, _), neg_out_list_ensemble, _ = neg_mask_sgd_ensemble(\n",
    "                    model, pos_img, neg_mask, c_repr, z=None, zgnn=None, wtarget=None, args=args,\n",
    "                    ensemble_size=ensemble_size,\n",
    "                    out_mode=\"min\",\n",
    "                )\n",
    "                neg_mask_ensemble_collection.append(neg_mask_ensemble)\n",
    "                neg_out_list_ensemble_collection.append(neg_out_list_ensemble)\n",
    "            neg_mask_ensemble_collection = tuple(torch.stack([neg_mask_ensemble_ele[k] for neg_mask_ensemble_ele in neg_mask_ensemble_collection], -1) for k in range(len(neg_mask_ensemble_collection[0])))\n",
    "            neg_out_list_ensemble_collection = np.stack(neg_out_list_ensemble_collection, -1)  # [n_steps, ensemble_size, batch_size, n_repr]\n",
    "            c_repr_energy = neg_out_list_ensemble_collection[-1].min(0)\n",
    "            c_repr_argsort = c_repr_energy.argsort(1)\n",
    "            c_repr_pred_list = []\n",
    "            for i, argsort in enumerate(c_repr_argsort):\n",
    "                c_repr_pred = {}\n",
    "                for k in range(len(args.concept_collection)):\n",
    "                    id_k = c_repr_argsort[i][k]\n",
    "                    c_repr_pred[args.concept_collection[id_k]] = c_repr_energy[i][id_k].item()\n",
    "                c_repr_pred_list.append(c_repr_pred)\n",
    "            data[\"both|c\"] = {\n",
    "                \"neg_mask_ensemble_collection\": neg_mask_ensemble_collection,\n",
    "                \"neg_out_list_ensemble_collection\": neg_out_list_ensemble_collection,\n",
    "                \"c_repr_pred\": c_repr_pred_list,\n",
    "            }\n",
    "        if isplot:\n",
    "            plot_data(model, data, args, plot_ensemble_mode=plot_ensemble_mode, w_type=w_type, plot_grey_scale=plot_grey_scale, topk=plot_topk)\n",
    "    else:\n",
    "        neg_out_list = []\n",
    "        neg_img, neg_id = neg_data\n",
    "        neg_img.requires_grad = True\n",
    "        noise = torch.randn(batch_size, args.in_channels, args.image_size, args.image_size, device=device)\n",
    "        for k in range(sample_step):\n",
    "            if \"noise\" not in locals():\n",
    "                noise = torch.randn(neg_img.shape[0], args.in_channels, args.image_size, args.image_size, device=device)\n",
    "\n",
    "            noise.normal_(0, args.lambd)\n",
    "            neg_img.data.add_(noise.data)\n",
    "\n",
    "            neg_out = model(neg_img, neg_id)\n",
    "            neg_out.sum().backward()\n",
    "            neg_out_list.append(to_np_array(neg_out))\n",
    "            neg_img.grad.data.clamp_(-0.01, 0.01)\n",
    "\n",
    "            neg_img.data.add_(neg_img.grad.data, alpha=-args.step_size)\n",
    "\n",
    "            neg_img.grad.detach_()\n",
    "            neg_img.grad.zero_()\n",
    "\n",
    "            neg_img.data.clamp_(0, 1)\n",
    "\n",
    "        data = {\n",
    "            \"neg_img\": neg_img,\n",
    "            \"neg_id\": neg_id,\n",
    "        }\n",
    "        neg_out_list = np.concatenate(neg_out_list, -1).T\n",
    "        if isplot:\n",
    "            if neg_img.shape[1] != 3:\n",
    "                for img, neg_ele in zip(neg_img.argmax(1), neg_out):\n",
    "                    print(\"loss={:.6f}\".format(to_np_array(neg_ele)))\n",
    "                    visualize_matrices([img])\n",
    "    return data\n",
    "\n",
    "\n",
    "def plot_data(model, data, args, plot_ensemble_mode=\"min\", w_type=\"image+mask\", topk=3, plot_grey_scale=False):\n",
    "    plt.figure(figsize=(12,6))\n",
    "    neg_out_list_ensemble = data[\"mask|c_repr\"][\"neg_out_list_ensemble\"]  # [n_steps, ensemble_size, batch_size]\n",
    "    for i in range(min(neg_out_list_ensemble.shape[-1], 6)):\n",
    "        for k in range(neg_out_list_ensemble.shape[1]):\n",
    "            plt.plot(neg_out_list_ensemble[:,k,i], c=COLOR_LIST[i], label=\"example_{}\".format(i) if k==0 else None, alpha=0.4)\n",
    "    plt.legend()\n",
    "    plt.show()\n",
    "\n",
    "    pos_img, pos_mask, pos_id, pos_repr = data[\"pos_img\"], data[\"pos_mask\"], data[\"pos_id\"], data[\"pos_repr\"]\n",
    "    concept_collection = data[\"concept_collection\"]\n",
    "    if topk == -1:\n",
    "        topk = len(concept_collection)\n",
    "    topk = min(topk, len(concept_collection))\n",
    "    length = len(pos_repr)\n",
    "    pos_out_last = to_np_array(model(pos_img, mask=pos_mask, c_repr=pos_repr)).squeeze()\n",
    "    if args.is_two_branch:\n",
    "        if pos_img[0].shape[1] == 3:\n",
    "            pos_img_core = torch.cat([pos_img[0].detach().to('cpu'), pos_img[1].detach().to('cpu')])\n",
    "        else:\n",
    "            pos_img_core = torch.cat([onehot_to_RGB(pos_img[0]), onehot_to_RGB(pos_img[1])])\n",
    "        pos_mask_core = torch.cat([pos_mask[0].detach().to('cpu').round(), pos_mask[1].detach().to('cpu').round()])\n",
    "        for i in range(0, length, 6):\n",
    "            print(\"\\n{} to {}:\".format(i, i+5))\n",
    "            print(\"positive images: input (up) and target (down)\")\n",
    "            visualize_matrices(pos_img[0][i:i+6].argmax(1), images_per_row=6)\n",
    "            visualize_matrices(pos_img[1][i:i+6].argmax(1), images_per_row=6)\n",
    "            print(\"positive mask: input obj mask (up) and target obj mask (down)\")\n",
    "            visualize_matrices(pos_mask[0][i:i+6,0].round() if \"mask\" in w_type else pos_mask[0][i:i+6].argmax(1), images_per_row=6, subtitles=[\"\\n\".join([\"{}: {:.4f}\".format(\"[{}]\".format(key) if pos_id[i+j]==key else key, item) for k, (key, item) in enumerate(Dict.items()) if k < topk]) for j, Dict in enumerate(data[\"c_repr|mask\"][\"c_repr_pred\"][i:i+6])])\n",
    "            visualize_matrices(pos_mask[1][i:i+6,0].round() if \"mask\" in w_type else pos_mask[1][i:i+6].argmax(1), images_per_row=6, subtitles=[\"c:\"+\"\\n\".join([\"{}: {:.4f}\".format(\"[{}]\".format(key) if pos_id[i+j]==key else key, item) for k, (key, item) in enumerate(Dict.items()) if k < topk]) for j, Dict in enumerate(data[\"both|c\"][\"c_repr_pred\"][i:i+6])])\n",
    "            if \"mask|c_repr\" in data:\n",
    "                if \"neg_out_argmin\" not in locals():\n",
    "                    neg_mask_ensemble = data[\"mask|c_repr\"][\"neg_mask_ensemble\"]\n",
    "                    neg_out_argmin = neg_out_list_ensemble[-1].argmin(0)  # neg_out_list_ensemble: [time_step, ensemble_size, batch_size]\n",
    "                if plot_ensemble_mode == \"min\":\n",
    "                    print(\"negative mask | c_repr: input obj mask (up) and target obj mask (down), with lowest-loss element in the ensemble:\")\n",
    "                    visualize_matrices([neg_mask_ensemble[0][neg_out_argmin[k],k,0].round() if \"mask\" in w_type else neg_mask_ensemble[0][neg_out_argmin[k],k].argmax(0) for k in range(i, i+6)], images_per_row=6, subtitles=[\"{}: {:.4f}\".format(pos_id[k], neg_out_list_ensemble[-1,neg_out_argmin[k],k]) for k in range(i,i+6)])\n",
    "                    visualize_matrices([neg_mask_ensemble[1][neg_out_argmin[k],k,0].round() if \"mask\" in w_type else neg_mask_ensemble[1][neg_out_argmin[k],k].argmax(0)  for k in range(i, i+6)], images_per_row=6)\n",
    "                    if plot_grey_scale:\n",
    "                        plot_matrices([neg_mask_ensemble[0][neg_out_argmin[k],k,0] for k in range(i, i+6)] if \"mask\" in w_type else neg_mask_ensemble[0][neg_out_argmin[k],k].argmax(0), scale_limit=(0,1), images_per_row=6)\n",
    "                        plot_matrices([neg_mask_ensemble[1][neg_out_argmin[k],k,0] for k in range(i, i+6)] if \"mask\" in w_type else neg_mask_ensemble[1][neg_out_argmin[k],k].argmax(0), scale_limit=(0,1), images_per_row=6)\n",
    "                elif plot_ensemble_mode == \"all\":\n",
    "                    for j in range(ensemble_size):\n",
    "                        is_best = j == neg_out_argmin\n",
    "                        if j == 0:\n",
    "                            print(\"negative mask | c_repr: {}th element in the ensemble for input obj mask (up) and target obj mask (down)\".format(j))\n",
    "                        else:\n",
    "                            print(\"                        {}th element in the ensemble for input obj mask (up) and target obj mask (down)\".format(j))\n",
    "                        visualize_matrices([neg_mask_ensemble[0][j,k,0].round() if \"mask\" in w_type else neg_mask_ensemble[0][j,k].argmax(0) for k in range(i, i+6)], images_per_row=6, subtitles=[\"{}: {:.4f}{}\".format(pos_id[k], neg_out_list_ensemble[-1,j,k], \", best\" if is_best[k] else \"\") for k in range(i,i+6)])\n",
    "                        visualize_matrices([neg_mask_ensemble[1][j,k,0].round() if \"mask\" in w_type else neg_mask_ensemble[1][j,k].argmax(0) for k in range(i, i+6)], images_per_row=6)\n",
    "                        if plot_grey_scale:\n",
    "                            plot_matrices([neg_mask_ensemble[0][j,k,0] for k in range(i, i+6)], scale_limit=(0,1), images_per_row=6)\n",
    "                            plot_matrices([neg_mask_ensemble[1][j,k,0] for k in range(i, i+6)], scale_limit=(0,1), images_per_row=6)\n",
    "                else:\n",
    "                    raise\n",
    "            if \"both|c\" in data:\n",
    "                if \"c_repr_pred_c\" not in locals():\n",
    "                    neg_mask_ensemble_collection = data[\"both|c\"]['neg_mask_ensemble_collection']\n",
    "                    neg_out_list_ensemble_collection = data[\"both|c\"][\"neg_out_list_ensemble_collection\"]\n",
    "                    c_repr_pred_c = data[\"both|c\"]['c_repr_pred']\n",
    "                    neg_out_argsort_c = neg_out_list_ensemble_collection[-1].min(0).argsort(-1)  # length batch_size, each indicating the argsort of concept id\n",
    "                for j in range(topk):\n",
    "                    if j == 0:\n",
    "                        print(\"negative mask | c: top {}th prediction for input obj mask (up) and target obj mask (down)\".format(j+1))\n",
    "                    else:\n",
    "                        print(\"                   top {}th prediction in the ensemble for input obj mask (up) and target obj mask (down)\".format(j+1))\n",
    "\n",
    "                    visualize_matrices([neg_mask_ensemble_collection[0][k,...,neg_out_argsort_c[k][j]].squeeze(0).round() \n",
    "                                        if \"mask\" in w_type else neg_mask_ensemble_collection[0][k,...,neg_out_argsort_c[k][j]].argmax(0)\n",
    "                                        for k in range(i, i+6)], images_per_row=6,\n",
    "                                       subtitles=[\"{}: {:.4f}\".format(\"[{}]\".format(concept_collection[neg_out_argsort_c[k][j]]) if concept_collection[neg_out_argsort_c[k][j]]==pos_id[k] else concept_collection[neg_out_argsort_c[k][j]], c_repr_pred_c[k][concept_collection[neg_out_argsort_c[k][j]]]) for k in range(i, i+6)]\n",
    "                                      )\n",
    "                    visualize_matrices([neg_mask_ensemble_collection[1][k,...,neg_out_argsort_c[k][j]].squeeze(0).round() \n",
    "                                        if \"mask\" in w_type else neg_mask_ensemble_collection[1][k,...,neg_out_argsort_c[k][j]].argmax(0)\n",
    "                                        for k in range(i, i+6)], images_per_row=6,\n",
    "                                       )\n",
    "            print()\n",
    "    else:\n",
    "        if pos_img.shape[1] == 3:\n",
    "            pos_img_core = pos_img.detach().to('cpu')\n",
    "        else:\n",
    "            pos_img_core = onehot_to_RGB(pos_img)\n",
    "        for i in range(0, length, 6):\n",
    "            print(\"\\n{} to {}:\".format(i, i+5))\n",
    "            print(\"positive images:\")\n",
    "            visualize_matrices(pos_img[i:i+6].argmax(1), images_per_row=6)\n",
    "            print(\"positive masks:\")\n",
    "            visualize_matrices(pos_mask[0][i:i+6,0].round() if \"mask\" in w_type else pos_mask[0][i:i+6].argmax(1), images_per_row=6, subtitles=[\"\\n\".join([\"{}: {:.4f}\".format(\"[{}]\".format(key) if pos_id[i+j]==key else key, item) for k, (key, item) in enumerate(Dict.items()) if k < topk]) + \"\\nc:\" + \"\\n\".join([\"{}: {:.4f}\".format(\"[{}]\".format(key) if pos_id[i+j]==key else key, item) for k, (key, item) in enumerate(data[\"both|c\"][\"c_repr_pred\"][i+j].items()) if k < topk]) for j, Dict in enumerate(data[\"c_repr|mask\"][\"c_repr_pred\"][i:i+6])])\n",
    "            if \"mask|c_repr\" in data:\n",
    "                if \"neg_out_argmin\" not in locals():\n",
    "                    neg_mask_ensemble = data[\"mask|c_repr\"][\"neg_mask_ensemble\"]\n",
    "                    neg_out_argmin = neg_out_list_ensemble[-1].argmin(0)\n",
    "                if plot_ensemble_mode == \"min\":\n",
    "                    print(\"negative mask | c_repr: for lowest-loss element in the ensemble:\")\n",
    "                    visualize_matrices([neg_mask_ensemble[0][neg_out_argmin[k],k,0].round()\n",
    "                                        if \"mask\" in w_type else neg_mask_ensemble[0][neg_out_argmin[k],k].argmax(0)\n",
    "                                        for k in range(i, i+6)], images_per_row=6, subtitles=[\"{}: {:.4f}\".format(pos_id[k], neg_out_list_ensemble[-1,neg_out_argmin[k],k]) for k in range(i,i+6)])\n",
    "                    if plot_grey_scale:\n",
    "                        plot_matrices([neg_mask_ensemble[0][neg_out_argmin[k],k,0] for k in range(i, i+6)], scale_limit=(0,1), images_per_row=6)\n",
    "                elif plot_ensemble_mode == \"all\":\n",
    "                    for j in range(ensemble_size):\n",
    "                        if j == 0:\n",
    "                            print(\"negative mask | c_repr: {}th element in the ensemble\".format(j))\n",
    "                        else:\n",
    "                            print(\"                        {}th element in the ensemble\".format(j))\n",
    "                        is_best = j == neg_out_argmin\n",
    "                        visualize_matrices([neg_mask_ensemble[0][j,k,0].round() for k in range(i, i+6)], images_per_row=6, subtitles=[\"{}: {:.4f}{}\".format(pos_id[k], neg_out_list_ensemble[-1,j,k], \", best\" if is_best[k] else \"\") for k in range(i,i+6)])\n",
    "                        if plot_grey_scale:\n",
    "                            plot_matrices([neg_mask_ensemble[0][j,k,0] for k in range(i, i+6)], scale_limit=(0,1), images_per_row=6)\n",
    "                else:\n",
    "                    raise\n",
    "            if \"both|c\" in data:\n",
    "                if \"c_repr_pred_c\" not in locals():\n",
    "                    neg_mask_ensemble_collection = data[\"both|c\"]['neg_mask_ensemble_collection']\n",
    "                    neg_out_list_ensemble_collection = data[\"both|c\"][\"neg_out_list_ensemble_collection\"]\n",
    "                    c_repr_pred_c = data[\"both|c\"]['c_repr_pred']\n",
    "                    neg_out_argsort_c = neg_out_list_ensemble_collection[-1].min(0).argsort(-1)  # length batch_size, each indicating the argsort of concept id\n",
    "                for j in range(topk):\n",
    "                    if j == 0:\n",
    "                        print(\"negative mask | c: top {}th prediction for the mask\".format(j+1))\n",
    "                    else:\n",
    "                        print(\"                   top {}th prediction for the mask\".format(j+1))\n",
    "\n",
    "                    visualize_matrices([neg_mask_ensemble_collection[0][k,...,neg_out_argsort_c[k][j]].squeeze(0).round() \n",
    "                                        if \"mask\" in w_type else neg_mask_ensemble_collection[0][k,...,neg_out_argsort_c[k][j]].argmax(0)\n",
    "                                        for k in range(i, i+6)], images_per_row=6,\n",
    "                                       subtitles=[\"{}: {:.4f}\".format(\"[{}]\".format(concept_collection[neg_out_argsort_c[k][j]]) if concept_collection[neg_out_argsort_c[k][j]]==pos_id[k] else concept_collection[neg_out_argsort_c[k][j]], c_repr_pred_c[k][concept_collection[neg_out_argsort_c[k][j]]]) for k in range(i, i+6)]\n",
    "                                      )\n",
    "            print()\n",
    "\n",
    "\n",
    "def plot_acc(dirname, filenames, acc_modes=None, filter_mode=None, prefix=None, suffix=None, is_plot_loss=True):\n",
    "    # Plot the accuracy for the a set of model given paths. You can specify which accuracy modes, prefix, and suffix to be used. \n",
    "    # If acc_modes is None, we use all accuracy modes\n",
    "    for filename in filenames: \n",
    "        path = dirname + filename\n",
    "        data_record = pickle.load(open(path, \"rb\"))\n",
    "        args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "        accuracy = data_record[\"acc\"]\n",
    "\n",
    "        table_acc_modes = {'Filenames': filenames}\n",
    "        if acc_modes == None:\n",
    "            acc_modes = accuracy.keys()\n",
    "        if filter_mode is None:\n",
    "            acc_modes = [acc_mode for acc_mode in acc_modes if (prefix is None or acc_mode.startswith(prefix)) and (suffix is None or acc_mode.endswith(suffix))]\n",
    "        elif filter_mode == \"standard\":\n",
    "            acc_modes = ['acc:mask|c_repr:val', 'acc:mask|c:val', 'acc:c_repr|mask:val', 'acc:c_repr|c:val']\n",
    "        else:\n",
    "            raise\n",
    "\n",
    "        print(filename)\n",
    "        plt.figure(figsize=(8,6))\n",
    "        best_dict = {}\n",
    "        for key in acc_modes:      \n",
    "            x = accuracy[\"epoch:val\"]\n",
    "            y = accuracy[key]\n",
    "            best_dict[key] = np.max(accuracy[key])\n",
    "            plt.plot(x, y, label=key)\n",
    "        acc_mean = np.array([accuracy[key] for key in acc_modes]).mean(0)\n",
    "        acc_mean_argmax = acc_mean.argmax()\n",
    "        plt.plot(x, acc_mean, label=\"mean_plot\")\n",
    "        plt.axvline(x=accuracy[\"epoch:val\"][acc_mean_argmax], color='k', linestyle='--')\n",
    "        if is_plot_loss:\n",
    "            if args.train_mode == \"cd\":\n",
    "                plt.plot(data_record[\"epoch\"], data_record[\"E:pos:train\"], label=\"E:pos:train\")\n",
    "                plt.plot(data_record[\"epoch\"], data_record[\"E:neg|c_repr:train\"], label=\"E:neg:train\")\n",
    "                plt.plot(data_record[\"epoch\"], data_record[\"E:neg_gen:train\"], label=\"E:neg_gen:train\")\n",
    "                if \"loss_kl\" in data_record:\n",
    "                    plt.plot(data_record[\"epoch\"], data_record[\"loss_kl\"], label=\"loss_kl\")\n",
    "                    plt.plot(data_record[\"epoch\"], data_record[\"loss_entropy_mask_mean\"], label=\"loss_entropy_mask_mean\")\n",
    "                    plt.plot(data_record[\"epoch\"], data_record[\"loss_entropy_repr_mean\"], label=\"loss_entropy_repr_mean\")\n",
    "            elif args.train_mode == \"sl\":\n",
    "                plt.plot(data_record[\"epoch\"], data_record[\"loss_mask:train\"], label=\"loss_mask:train\")\n",
    "                plt.plot(data_record[\"epoch\"], data_record[\"loss_repr:train\"], label=\"loss_repr:train\")\n",
    "                plt.plot(data_record[\"epoch\"], data_record[\"loss:train\"], label=\"loss:train\")\n",
    "            else:\n",
    "                raise\n",
    "\n",
    "        plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))\n",
    "        title = \"Best:\"\n",
    "        for key in best_dict:\n",
    "            title += \"  {}={:.4f}\".format(key[4:-4], best_dict[key], end=\"\")\n",
    "        plt.title(title)\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "def analyze(\n",
    "    dirname,\n",
    "    filename,\n",
    "    init=\"gaussian\", # \"mask-gaussian\"\n",
    "    batch_size=12,\n",
    "    load_epoch=-1,\n",
    "    ensemble_size=8,\n",
    "    plot_ensemble_mode=\"min\",\n",
    "    plot_grey_scale=False,\n",
    "    analysis_result=None,\n",
    "    # New settings:\n",
    "    n_examples=1000,\n",
    "    sample_step=120,\n",
    "    lambd=-1,\n",
    "    seed=2,\n",
    "    CONCEPTS=None,\n",
    "    OPERATORS=None,\n",
    "    isplot=True,\n",
    "):\n",
    "    \"\"\"Get test_acc, generate samples, and visualize.\"\"\"\n",
    "    print_banner(filename)\n",
    "    try:\n",
    "        data_record = pickle.load(open(dirname + filename, \"rb\"))\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        return None, None, None\n",
    "    if isplot:\n",
    "        plot_loss(data_record)\n",
    "    args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "    pp.pprint(args.__dict__)\n",
    "    if 'concept_embeddings' in data_record:\n",
    "        test_concept_embedding(CONCEPTS, OPERATORS, data_record['concept_embeddings'])\n",
    "\n",
    "    if n_examples is not None:\n",
    "        args.n_examples = n_examples\n",
    "    args.gpuid = \"False\"\n",
    "    if lambd != -1:\n",
    "        args.lambd = lambd\n",
    "    device = \"cpu\"\n",
    "    if load_epoch < 0:\n",
    "        load_epoch += data_record[\"epoch\"][-1]\n",
    "    id = int(np.round(load_epoch / args.save_interval))\n",
    "    if id > len(data_record[\"model_dict\"]):\n",
    "        print(\"{} has only {} epochs. Continue\".format(filename, data_record[\"epoch\"][-1]))\n",
    "        return None, None, None\n",
    "\n",
    "    if analysis_result is None:\n",
    "        set_seed(seed=seed)\n",
    "        dataset, args = get_dataset(args)\n",
    "        dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=True, collate_fn=Batch(is_collate_tuple=True).collate())\n",
    "        model = load_model_energy(data_record[\"model_dict\"][id], device=device)\n",
    "        val_acc_dict = test_acc(model, args, dataloader, device, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS, suffix=\":val\")\n",
    "        print(\"\\nacc and energy:\")\n",
    "        pp.pprint(val_acc_dict)\n",
    "        print(\"\\nmodel at epoch {}:\".format(id * args.save_interval))\n",
    "        # Perform gradient descent to find a low energy mask:\n",
    "        data, neg_out_list_ensemble = get_model_samples(model, args, dataset=dataset, init=init, sample_step=sample_step, batch_size=batch_size, ensemble_size=ensemble_size, plot_ensemble_mode=plot_ensemble_mode, plot_grey_scale=plot_grey_scale, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS, isplot=isplot)\n",
    "    else:\n",
    "        print(\"Loading saved analysis:\")\n",
    "        model = load_model_energy(data_record[\"model_dict\"][id], device=device)\n",
    "        data = analysis_result[\"data\"]\n",
    "        neg_out_list_ensemble = analysis_result[\"neg_out_list_ensemble\"]\n",
    "        val_acc_dict = analysis_result[\"val_acc_dict\"]\n",
    "        print(\"\\nacc and energy:\")\n",
    "        pp.pprint(val_acc_dict)\n",
    "        print(\"\\nmodel at epoch {}:\".format(id * args.save_interval))\n",
    "        if isplot:\n",
    "            plot_data(model, data, neg_out_list_ensemble, args, plot_ensemble_mode=plot_ensemble_mode, plot_grey_scale=plot_grey_scale)\n",
    "        dataset = None\n",
    "    print()\n",
    "    data_record.pop(\"model_dict\")\n",
    "    info = {\n",
    "        \"args\": args,\n",
    "        \"data_record\": data_record,\n",
    "        \"data\": data,\n",
    "        \"neg_out_list_ensemble\": neg_out_list_ensemble,\n",
    "        \"val_acc_dict\": val_acc_dict,\n",
    "    }\n",
    "    return model, dataset, info\n",
    "\n",
    "\n",
    "def get_useful_keys(df, args_keys):\n",
    "    useful_keys = []\n",
    "    for key in sorted(args_keys):\n",
    "        try:\n",
    "            df_mean = df.groupby(by=key).mean()\n",
    "        except:\n",
    "            continue\n",
    "        if len(df_mean) == 1:\n",
    "            continue\n",
    "        useful_keys.append(key)\n",
    "    if \"gpuid\" in useful_keys:\n",
    "        useful_keys.remove(\"gpuid\")\n",
    "    if \"is_two_branch\" in useful_keys:\n",
    "        useful_keys.remove(\"is_two_branch\")\n",
    "    return useful_keys\n",
    "\n",
    "\n",
    "def print_info(args, keys=None):\n",
    "    if keys is None:\n",
    "        keys = [\"dataset\", \"canvas_size\", \"neg_mode_coef\", \"ebm_target_mode\", \"step_size_repr\", \"c_repr_mode\", \"c_repr_first\", \"channel_base\", \"n_examples\", \"aggr_mode\", \"kl_coef\", \"entropy_coef_mask\", \"entropy_coef_repr\", \"is_spec_norm\", \"color_avail\",\"id\"]\n",
    "    for key in keys:\n",
    "        print(\"{}: {}\".format(key, getattr(args, key)))\n",
    "\n",
    "\n",
    "def get_update_acc(dirname, suffix=\"\"):\n",
    "    \"\"\"Get Pandas dataframe analysing acc.\"\"\"\n",
    "    filenames = filter_filename(dirname, include=\".p\")\n",
    "    df_dict_list = []\n",
    "    for filename in filenames:\n",
    "        data_record = pickle.load(open(dirname + filename, \"rb\"))\n",
    "        args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "        df_dict = {}\n",
    "        acc_dict = data_record[\"acc\"]\n",
    "        for key in acc_dict:\n",
    "            best_id_acc = np.argmax(acc_dict['acc:mean:val'])\n",
    "            acc_repr = np.array([acc_dict[key] for key in acc_dict if key.startswith(\"acc:c_repr\")]).mean(0)\n",
    "            acc_mask = np.array([acc_dict[key] for key in acc_dict if key.startswith(\"acc:mask\") and \"_0\" not in key and \"_1\" not in key]).mean(0)\n",
    "            best_id_acc_repr = np.argmax(acc_repr)\n",
    "            best_id_acc_mask = np.argmax(acc_mask)\n",
    "            df_dict[\"epoch\"] = acc_dict[\"epoch:val\"][-1]\n",
    "            df_dict[\"best_epoch:acc\"] = acc_dict[\"epoch:val\"][best_id_acc]\n",
    "            df_dict[\"best_epoch:acc_repr\"] = acc_dict[\"epoch:val\"][best_id_acc_repr]\n",
    "            df_dict[\"best_epoch:acc_mask\"] = acc_dict[\"epoch:val\"][best_id_acc_mask]\n",
    "            df_dict[\"acc:c_repr:mean:val\"] = np.max(acc_repr)\n",
    "            df_dict[\"acc:mask:mean:val\"] = np.max(acc_mask)\n",
    "\n",
    "            #print(df_dict[\"best_epoch:acc_mask\"])\n",
    "        best_id_model = int(df_dict[\"best_epoch:acc_mask\"] / 20)\n",
    "        decice= \"cuda:0\"\n",
    "        model = load_model_energy(data_record[\"model_dict\"][best_id_model], device=device)\n",
    "        set_seed(seed=2)\n",
    "        args.n_examples = 500\n",
    "        dataset, _ = get_dataset(args)\n",
    "        dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=False, collate_fn=Batch(is_collate_tuple=True).collate())\n",
    "        decice= \"cuda:0\"\n",
    "        update_val_acc_dict = test_acc(model, args, dataloader, device, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS, suffix=suffix)\n",
    "        print(filename)\n",
    "        print(\"epoch: \", best_id_model)\n",
    "        print(update_val_acc_dict)\n",
    "        print(\"\\n\")\n",
    "\n",
    "\n",
    "def get_df(dirname):\n",
    "    \"\"\"Get Pandas dataframe analysising acc.\"\"\"\n",
    "    filenames = filter_filename(dirname, include=\".p\")\n",
    "    df_dict_list = []\n",
    "    for filename in filenames:\n",
    "        try:\n",
    "            data_record = pickle.load(open(dirname + filename, \"rb\"))\n",
    "        except Exception as e:\n",
    "            print(\"Error {} from file {}\".format(e, filename))\n",
    "        args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "\n",
    "        df_dict = {}\n",
    "        acc_dict = data_record[\"acc\"]\n",
    "        for key in acc_dict:\n",
    "            best_id_acc = np.argmax(acc_dict['acc:mean:val'])\n",
    "            acc_repr = np.array([acc_dict[key] for key in acc_dict if key.startswith(\"acc:c_repr\")]).mean(0)\n",
    "            acc_mask = np.array([acc_dict[key] for key in acc_dict if key.startswith(\"acc:mask\") and \"_0\" not in key and \"_1\" not in key]).mean(0)\n",
    "            best_id_acc_repr = np.argmax(acc_repr)\n",
    "            best_id_acc_mask = np.argmax(acc_mask)\n",
    "            df_dict[\"epoch\"] = acc_dict[\"epoch:val\"][-1]\n",
    "            df_dict[\"best_epoch:acc\"] = acc_dict[\"epoch:val\"][best_id_acc]\n",
    "            df_dict[\"best_epoch:acc_repr\"] = acc_dict[\"epoch:val\"][best_id_acc_repr]\n",
    "            df_dict[\"best_epoch:acc_mask\"] = acc_dict[\"epoch:val\"][best_id_acc_mask]\n",
    "            df_dict[\"acc:c_repr:mean:val\"] = np.max(acc_repr)\n",
    "            df_dict[\"acc:mask:mean:val\"] = np.max(acc_mask)\n",
    "\n",
    "            if key == \"epoch:val\":\n",
    "                pass\n",
    "            elif key.startswith(\"acc:\"):\n",
    "                df_dict[key] = np.max(acc_dict[key])\n",
    "            elif key.startswith(\"E:\"):\n",
    "                df_dict[key] = acc_dict[key][best_id_acc]\n",
    "            else:\n",
    "                raise\n",
    "        df_dict.update(args.__dict__)\n",
    "        df_dict[\"filename\"] = filename\n",
    "        df_dict_list.append(df_dict)\n",
    "\n",
    "    df = pd.DataFrame(df_dict_list)\n",
    "    info = {\"args_keys\": list(args.__dict__.keys()),\n",
    "            \"acc_keys\": [key for key in acc_dict.keys() if \"_0\" not in key and \"_1\" not in key],\n",
    "           }\n",
    "    return df, info\n",
    "\n",
    "\n",
    "def get_bidirectional_graph(graph):\n",
    "    Dict = {\"IsInside\": \"IsEnclosed\",\n",
    "            \"IsEnclosed\": \"IsInside\",\n",
    "           }\n",
    "    graph_new = []\n",
    "    # Standardize:\n",
    "    def standardize_graph(graph):\n",
    "        new_graph = []\n",
    "        for ele in graph:\n",
    "            if isinstance(ele[0], list):\n",
    "                ele = (tuple(ele[0]), ele[1])\n",
    "            new_graph.append(ele)\n",
    "        return new_graph\n",
    "    new_graph = standardize_graph(graph)\n",
    "    graph_add = []\n",
    "\n",
    "    for ele in new_graph:\n",
    "        if isinstance(ele[0], tuple):\n",
    "            reverse_ele = ((ele[0][1], ele[0][0]), ele[1])\n",
    "            if reverse_ele not in new_graph:\n",
    "                graph_add.append(reverse_ele)\n",
    "    return new_graph + graph_add"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebbec290-7d45-4798-803a-47dbfeb957e3",
   "metadata": {},
   "source": [
    "### 1.1 Helper functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52b2dd72-108f-421e-aa43-5688462b5086",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_model_hash(\n",
    "    hash_str,\n",
    "    is_dataset=False,\n",
    "    is_update_repr=False,\n",
    "    isplot=1,\n",
    "    mutual_exclusive_coef=None,\n",
    "    return_args=False,\n",
    "    update_keys=None,\n",
    "    load_epoch=\"best\",\n",
    "):\n",
    "    try:\n",
    "        load_epoch = eval(load_epoch)\n",
    "    except:\n",
    "        pass\n",
    "    subfolders = filter_filename(EXP_PATH, exclude=\".\")\n",
    "    is_found = False\n",
    "    for subfolder in subfolders:\n",
    "        if \"evaluation\" in subfolder:\n",
    "            continue\n",
    "        filenames = filter_filename(f\"{EXP_PATH}/{subfolder}/\", include=[hash_str, \".p\"])\n",
    "        if len(filenames) == 1:\n",
    "            is_found = True\n",
    "            break\n",
    "    assert is_found, f\"Did not find the experiment with hash_str '{hash_str}' under the subfolders of './{EXP_PATH}/'. Please check if the hash_str is correct.\"\n",
    "    filename = filenames[0]\n",
    "    data_record = pickle.load(open(f\"{EXP_PATH}/{subfolder}/{filename}\", \"rb\"))\n",
    "    if is_update_repr:\n",
    "        update_CONCEPTS_OPERATORS(CONCEPTS, OPERATORS, data_record[\"concept_embeddings\"][-1], update_keys=update_keys)\n",
    "        test_concept_embedding(CONCEPTS, OPERATORS, data_record[\"concept_embeddings\"][-1], raise_warnings_only=True, checked_keys=update_keys)\n",
    "    else:\n",
    "        test_concept_embedding(CONCEPTS, OPERATORS, data_record[\"concept_embeddings\"][-1], checked_keys=update_keys)\n",
    "    args = init_args(update_default_hyperparam(data_record[\"args\"]))\n",
    "    if isplot >= 1:\n",
    "        plot_acc(dirname, [filename], filter_mode=\"standard\")\n",
    "    model, best_model_id = load_best_model(data_record,\n",
    "                            keys=[\"mask|c_repr\", \"mask|c\", \"c_repr|mask\", \"c_repr|c\"],\n",
    "                            load_epoch=load_epoch,\n",
    "                            return_id=True,\n",
    "                           )\n",
    "    model.to(device)\n",
    "\n",
    "    info = {\"load_id\": best_model_id}\n",
    "    if return_args:\n",
    "        info[\"args\"] = args\n",
    "    if is_dataset:\n",
    "        args.n_examples = 500\n",
    "        args.seed = 2\n",
    "        print(\"canvas_size: {}, color_avail: {}\".format(args.canvas_size, args.color_avail))\n",
    "        dataset, args = get_dataset(args, is_load=True)\n",
    "        info[\"dataset\"] = dataset\n",
    "\n",
    "        if isplot >= 2:\n",
    "            init = \"mask-input\"\n",
    "            sample_step = 150\n",
    "            batch_size = 36\n",
    "            ensemble_size = 16\n",
    "            plot_grey_scale = False\n",
    "            plot_ensemble_mode = \"min\"\n",
    "            args.lambd_start = 0.1\n",
    "            if mutual_exclusive_coef is None:\n",
    "                mutual_exclusive_coef = args.mutual_exclusive_coef\n",
    "            args.mutual_exclusive_coef = mutual_exclusive_coef\n",
    "            data = get_model_samples(model, args, dataset=dataset, init=init, sample_step=sample_step, batch_size=batch_size, ensemble_size=ensemble_size, plot_ensemble_mode=plot_ensemble_mode, plot_grey_scale=plot_grey_scale, CONCEPTS=CONCEPTS, OPERATORS=OPERATORS, device=device, isplot=True)\n",
    "    return model, info\n",
    "\n",
    "\n",
    "# Build an empty selector:\n",
    "def get_empty_selector(ebm_dict, CONCEPTS, OPERATORS):\n",
    "    selector = Concept_Pattern(\n",
    "        name=None,\n",
    "        value=Placeholder(Tensor(dtype=\"cat\", range=range(10))),\n",
    "        attr={},\n",
    "        is_all_obj=True,\n",
    "        is_ebm=True,\n",
    "        is_selector_gnn=False,\n",
    "        is_default_ebm=False,\n",
    "        ebm_dict=ebm_dict,\n",
    "        CONCEPTS=CONCEPTS,\n",
    "        OPERATORS=OPERATORS,\n",
    "        device=device,\n",
    "        cache_forward=False,\n",
    "        in_channels=10,\n",
    "        z_mode=\"None\",\n",
    "        w_type=\"image+mask\",\n",
    "        mask_mode=\"concat\",\n",
    "        aggr_mode=\"max\",\n",
    "        pos_embed_mode=\"None\",\n",
    "        is_ebm_share_param=True,\n",
    "        is_relation_z=False,\n",
    "        img_dims=2,\n",
    "        is_spec_norm=True,\n",
    "        act_name=\"leakyrelu0.2\",\n",
    "        normalization_type=\"None\",\n",
    "    )\n",
    "    return selector\n",
    "\n",
    "\n",
    "def get_c_repr(c_str, num, device, mode):\n",
    "    if mode == \"concept\":\n",
    "        c_repr = CONCEPTS[c_str].get_node_repr().detach()[None].repeat_interleave(num, 0).to(device)\n",
    "    elif mode == \"relation\":\n",
    "        c_repr = OPERATORS[c_str].get_node_repr().detach()[None].repeat_interleave(num, 0).to(device)\n",
    "    else:\n",
    "        raise\n",
    "    return c_repr\n",
    "\n",
    "\n",
    "def get_graph_info(info):\n",
    "    # Get all objects and their concepts:\n",
    "    graph_info = {\"obj_type\": {}, \"relations\": {}}\n",
    "    obj_masks = info[\"obj_masks\"]\n",
    "    graph_info[\"obj_masks\"] = obj_masks\n",
    "    for key, mask in obj_masks.items():\n",
    "        obj_type = classify_concept(mask)\n",
    "        graph_info[\"obj_type\"][key] = obj_type\n",
    "    # Get relations:\n",
    "    for key, relation in info[\"obj_full_relations\"].items():\n",
    "        relations_valid = {\"SameShape\", \"SameColor\", \"IsInside\"}\n",
    "        relation_ele_valid = set(relation).intersection(relations_valid)\n",
    "        if len(relation_ele_valid) == 0:\n",
    "            continue\n",
    "        if len(relation_ele_valid) == 2:\n",
    "            pass\n",
    "        graph_info[\"relations\"][key] = list(relation_ele_valid)\n",
    "    graph_info[\"structure\"] = info[\"structure\"]\n",
    "    return graph_info\n",
    "\n",
    "\n",
    "def get_concept_from_graph_info(graph_info):\n",
    "    attr_dict = {}\n",
    "    for key, obj_type in graph_info[\"obj_type\"].items():\n",
    "        attr_dict[\"{}\".format(key, obj_type)] = Placeholder(obj_type, value=graph_info[\"obj_masks\"][key])\n",
    "\n",
    "    concept = Concept(\n",
    "        name=\"concept\",\n",
    "        value=Placeholder(Tensor(dtype=\"cat\", range=range(10))),\n",
    "        attr=attr_dict,\n",
    "    )\n",
    "\n",
    "    for key, relations in graph_info[\"relations\"].items():\n",
    "        for relation in relations:\n",
    "            concept.add_relation_manual(relation, key[0], key[1])\n",
    "    return concept\n",
    "\n",
    "\n",
    "def get_query_concept_pattern(query_type, graph_info):\n",
    "    def filter_relations(relations):\n",
    "        relations = deepcopy(relations)\n",
    "        keys_to_pop = []\n",
    "        for key, relation in relations.items():\n",
    "            if (key[1], key[0]) in relations and relations[(key[1], key[0])] == relation and key not in keys_to_pop:\n",
    "                keys_to_pop.append((key[1], key[0]))\n",
    "        for key in keys_to_pop:\n",
    "            relations.pop(key)\n",
    "        return relations\n",
    "\n",
    "    query_type = query_type.split(\"-\")[-1]\n",
    "    obj_names = list(graph_info[\"obj_type\"])\n",
    "#     graph_info = deepcopy(graph_info)\n",
    "#     graph_info[\"relations\"] = filter_relations(graph_info[\"relations\"])\n",
    "    concept = get_concept_from_graph_info(graph_info)\n",
    "\n",
    "    if query_type == \"1c\":\n",
    "        obj_type_avail = list(graph_info[\"obj_type\"].values())\n",
    "        obj_type_query = np.random.choice(obj_type_avail)\n",
    "        query_key_candidates = [key for key, obj_type in graph_info[\"obj_type\"].items() if obj_type == obj_type_query]\n",
    "        refer_masks = torch.stack([graph_info[\"obj_masks\"][key] for key in query_key_candidates])\n",
    "        query = {\n",
    "            \"graph\": [(0, obj_type_query)],\n",
    "            \"refer\": 0,\n",
    "        }\n",
    "    elif query_type == \"2a\":\n",
    "        # [concept, (0, 1)/(1, 0), (refer)]\n",
    "        assert len(graph_info[\"relations\"]) > 0\n",
    "        relation_id_selected = np.random.choice(len(graph_info[\"relations\"]))\n",
    "        relation_key_selected = list(graph_info[\"relations\"])[relation_id_selected]\n",
    "        key0, key1 = relation_key_selected\n",
    "        pivot_node_id = np.random.choice(2)\n",
    "        refer_node_id = 1 - pivot_node_id\n",
    "        pivot_node_key = relation_key_selected[pivot_node_id]\n",
    "        c_pattern = Concept_Pattern(\n",
    "            name=None,\n",
    "            value=Placeholder(Tensor(dtype=\"cat\", range=range(10))),\n",
    "            attr={key0: Placeholder(graph_info[\"obj_type\"][key0]),\n",
    "                  key1: Placeholder(graph_info[\"obj_type\"][key1]),\n",
    "                 },\n",
    "            re={relation_key_selected: graph_info[\"relations\"][relation_key_selected][0]},\n",
    "            refer_node_names=[relation_key_selected[refer_node_id]],\n",
    "        )\n",
    "        refer_masks = concept.get_refer_nodes(c_pattern, is_match_node=True)\n",
    "        refer_masks = {key: value.get_node_value() for key, value in refer_masks.items()}\n",
    "        query = {\"graph\": [((0,1),graph_info[\"relations\"][relation_key_selected][0]), (pivot_node_id, graph_info[\"obj_type\"][pivot_node_key])],\n",
    "                 \"refer\": refer_node_id}\n",
    "        refer_masks = torch.stack(list(refer_masks.values()))\n",
    "    elif query_type == \"3a\":\n",
    "        # [concept, (0, 1)/(1, 0), (1,2)/(2,1), (refer)]\n",
    "        if graph_info[\"structure\"] == ['pivot:Rect', (1, 0, 'IsInside'), '(concept)', (1, 2), '(refer)']:\n",
    "            query = {\n",
    "                \"graph\": [(0, \"Rect\"), ((1,0), \"IsInside\"), ((1,2), graph_info[\"relations\"][(obj_names[1], obj_names[2])][0])],\n",
    "                \"refer\": 2,\n",
    "            }\n",
    "            refer_masks = graph_info[\"obj_masks\"][obj_names[2]][None]\n",
    "        elif graph_info[\"structure\"] == [\"pivot\", (0,1), \"(concept)\", (1,2), \"(refer)\"]:\n",
    "            query = {\n",
    "                \"graph\": [(0, graph_info[\"obj_type\"][obj_names[0]]), ((0,1), graph_info[\"relations\"][(obj_names[0], obj_names[1])][0]), ((1,2), graph_info[\"relations\"][(obj_names[1], obj_names[2])][0])],\n",
    "                \"refer\": 2,\n",
    "            }\n",
    "            refer_masks = graph_info[\"obj_masks\"][obj_names[2]][None]\n",
    "        else:\n",
    "            return None, None\n",
    "    elif query_type == \"3b\":\n",
    "        if graph_info[\"structure\"] == [\"pivot\", \"pivot\", (0,2), (1,2), \"(refer)\"]:\n",
    "            query = {\n",
    "                \"graph\": [(0, graph_info[\"obj_type\"][obj_names[0]]),\n",
    "                          (1, graph_info[\"obj_type\"][obj_names[1]]),\n",
    "                          ((0,2), graph_info[\"relations\"][(obj_names[0], obj_names[2])][0]),\n",
    "                          ((1,2), graph_info[\"relations\"][(obj_names[1], obj_names[2])][0]),\n",
    "                         ],\n",
    "                \"refer\": 2,\n",
    "            }\n",
    "            refer_masks = graph_info[\"obj_masks\"][obj_names[2]][None]\n",
    "        else:\n",
    "            return None, None\n",
    "    return query, refer_masks\n",
    "\n",
    "\n",
    "def get_all_queries(dataset, query_type, seed=2, isplot=False):\n",
    "    set_seed(seed)\n",
    "    query_all_dict = {}\n",
    "    query_type = query_type.split(\"-\")[-1]\n",
    "    for i in range(len(dataset)):\n",
    "        input, target, infos = get_inputs_targets_EBM(dataset[i])\n",
    "        input = torch.stack(list(input[0].values())).to(device)\n",
    "        for j in range(len(input)):\n",
    "            graph_info = get_graph_info(infos[j])\n",
    "            query, refer_masks = get_query_concept_pattern(query_type, graph_info)\n",
    "            if query is not None:\n",
    "                record_data(query_all_dict, [input[j:j+1], query, refer_masks], [\"input\", \"query\", \"refer_masks\"])\n",
    "                if isplot:\n",
    "                    visualize_matrices(input[j:j+1].argmax(1))\n",
    "                    print(query)\n",
    "                    plot_matrices(refer_masks)\n",
    "                    print(\"\\n\\n\")\n",
    "    return query_all_dict\n",
    "\n",
    "\n",
    "def get_selector_from_graph(graph, ebm_dict, CONCEPTS, OPERATORS):\n",
    "    def get_is_exist(selector, node_name):\n",
    "        return node_name in [ele.split(\":\")[0] for ele in list(selector.nodes)]\n",
    "    selector = get_empty_selector(ebm_dict, CONCEPTS, OPERATORS)\n",
    "    for i, item in enumerate(graph):\n",
    "        if isinstance(item[0], tuple):\n",
    "            obj0 = \"obj_{}\".format(item[0][0]) if get_is_exist(selector, \"obj_{}\".format(item[0][0])) else None\n",
    "            obj1 = \"obj_{}\".format(item[0][1]) if get_is_exist(selector, \"obj_{}\".format(item[0][1])) else None\n",
    "            if obj0 is not None and obj1 is None:\n",
    "                selector.add_relation_manual(item[1], obj0)\n",
    "            elif obj0 is None and obj1 is not None:\n",
    "                selector.add_relation_manual(item[1], None, obj1)\n",
    "            elif obj0 is None and obj1 is None:\n",
    "                selector.add_relation_manual(item[1])\n",
    "            else:\n",
    "                assert obj0 is not None and obj1 is not None\n",
    "                selector.add_relation_manual(item[1], obj0, obj1)\n",
    "        else:\n",
    "            selector.add_obj(CONCEPTS[item[1]], add_obj_name=\"obj_{}\".format(item[0]))\n",
    "    if \"refer\" in graph:\n",
    "        selector.set_refer_nodes(\"obj_{}\".format(graph[\"refer\"]))\n",
    "    return selector\n",
    "\n",
    "\n",
    "def get_selector_for_parsing(keys_dict, ebm_dict, CONCEPTS, OPERATORS):\n",
    "    \"\"\"E.g. keys_dict = {\"Line\": 3, \"SameShape\": 4}\"\"\"\n",
    "    selector = get_empty_selector(ebm_dict, CONCEPTS, OPERATORS)\n",
    "    for key, count in keys_dict.items():\n",
    "        if key in CONCEPTS:\n",
    "            for i in range(count):\n",
    "                selector.add_obj(CONCEPTS[key])\n",
    "        elif key in OPERATORS:\n",
    "            for i in range(count):\n",
    "                selector.add_relation_manual(key)\n",
    "        else:\n",
    "            raise\n",
    "    return selector\n",
    "\n",
    "\n",
    "def get_concept_graph(c_type, is_new_vertical=True, is_bidirectional_re=False, is_concept=True, is_relation=True):\n",
    "    if is_new_vertical:\n",
    "        if not is_bidirectional_re:\n",
    "            graph_gt_dict = {\n",
    "                \"Line\": [\n",
    "                    (0, \"Line\"),\n",
    "                ],\n",
    "                \"2-Line\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                ],\n",
    "                \"3-Line\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                ],\n",
    "                \"Parallel\": [\n",
    "                    ((0, 1), \"Parallel\"),\n",
    "                ],\n",
    "                \"VerticalMid\": [\n",
    "                    ((0, 1), \"VerticalMid\"),\n",
    "                ],\n",
    "                \"VerticalEdge\": [\n",
    "                    ((0, 1), \"VerticalEdge\"),\n",
    "                ],\n",
    "                \"Eshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalMid'),\n",
    "                    ((0, 3), 'VerticalEdge'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((1, 3), 'Parallel'),\n",
    "                    ((2, 3), 'Parallel'),\n",
    "                ],\n",
    "                \"Fshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalMid'),\n",
    "                    ((1, 2), 'Parallel'),  \n",
    "                ],\n",
    "                \"Lshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                ],\n",
    "                \"Tshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    ((0, 1), 'VerticalMid'),\n",
    "                ],\n",
    "                \"Cshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalEdge'),\n",
    "                    ((1, 2), 'Parallel'),  \n",
    "                ],\n",
    "                \"Ashape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalMid'),\n",
    "                    ((0, 3), 'Parallel'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((1, 3), 'VerticalEdge'),\n",
    "                    ((2, 3), 'VerticalMid'),\n",
    "                ],\n",
    "                \"Hshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'VerticalMid'),\n",
    "                    ((1, 2), 'VerticalMid'),\n",
    "                    ((0, 2), 'Parallel'),\n",
    "                ],\n",
    "                \"Rect\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((1, 2), 'VerticalEdge'),\n",
    "                    ((2, 3), 'VerticalEdge'),\n",
    "                    ((0, 3), 'VerticalEdge'),\n",
    "                    ((0, 2), 'Parallel'),\n",
    "                    ((1, 3), 'Parallel'),\n",
    "                ],\n",
    "                \"RectE1a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                    ((1, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 1), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE1b\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                    ((1, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 1), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE2a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE2b\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE3a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((2, 0), \"IsInside\"),\n",
    "                    ((0, 2), \"IsEnclosed\"),\n",
    "                ],\n",
    "                \"RectE3b\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((2, 0), \"IsInside\"),\n",
    "                    ((0, 2), \"IsEnclosed\"),\n",
    "                ],\n",
    "                \"RectEconcept\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                ],\n",
    "                \"Graph1\": [\n",
    "                    (0, \"Red\"),\n",
    "                    ((0,1), \"SameColor\"),\n",
    "                    ((1,2), \"SameShape\"),\n",
    "                ],\n",
    "                \"Graph2\": [\n",
    "                    (0, \"Large\"),\n",
    "                    ((0,1), \"SameSize\"),\n",
    "                    ((0,2), \"SameColor\"),   \n",
    "                ],\n",
    "                \"Graph3\": [\n",
    "                    (0, \"Cube\"),\n",
    "                    ((0,1), \"SameShape\"),\n",
    "                    ((1,2), \"SameSize\"),\n",
    "                ],\n",
    "            }\n",
    "        else:\n",
    "            graph_gt_dict = {\n",
    "                \"Line\": [\n",
    "                    (0, \"Line\"),\n",
    "                ],\n",
    "                \"2-Line\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                ],\n",
    "                \"3-Line\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                ],\n",
    "                \"Parallel\": [\n",
    "                    ((0, 1), \"Parallel\"),\n",
    "                    ((1, 0), \"Parallel\"),\n",
    "                ],\n",
    "                \"VerticalMid\": [\n",
    "                    ((0, 1), \"VerticalMid\"),\n",
    "                    ((1, 0), \"VerticalMid\"),\n",
    "                ],\n",
    "                \"VerticalEdge\": [\n",
    "                    ((0, 1), \"VerticalEdge\"),\n",
    "                    ((1, 0), \"VerticalEdge\"),\n",
    "                ],\n",
    "                \"Eshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((1, 0), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalMid'),\n",
    "                    ((2, 0), 'VerticalMid'),\n",
    "                    ((0, 3), 'VerticalEdge'),\n",
    "                    ((3, 0), 'VerticalEdge'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((2, 1), 'Parallel'),\n",
    "                    ((1, 3), 'Parallel'),\n",
    "                    ((3, 1), 'Parallel'),\n",
    "                    ((2, 3), 'Parallel'),\n",
    "                    ((3, 2), 'Parallel'),\n",
    "                ],\n",
    "                \"Fshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((1, 0), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalMid'),\n",
    "                    ((2, 0), 'VerticalMid'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((2, 1), 'Parallel'),\n",
    "                ],\n",
    "                \"Lshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((1, 0), 'VerticalEdge'),\n",
    "                ],\n",
    "                \"Tshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    ((0, 1), 'VerticalMid'),\n",
    "                    ((1, 0), 'VerticalMid'),\n",
    "                ],\n",
    "                \"Cshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((1, 0), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalEdge'),\n",
    "                    ((2, 0), 'VerticalEdge'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((2, 1), 'Parallel'),\n",
    "                ],\n",
    "                \"Ashape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((1, 0), 'VerticalEdge'),\n",
    "                    ((0, 2), 'VerticalMid'),\n",
    "                    ((2, 0), 'VerticalMid'),\n",
    "                    ((0, 3), 'Parallel'),\n",
    "                    ((3, 0), 'Parallel'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((2, 1), 'Parallel'),\n",
    "                    ((1, 3), 'VerticalEdge'),\n",
    "                    ((3, 1), 'VerticalEdge'),\n",
    "                    ((2, 3), 'VerticalMid'),\n",
    "                    ((3, 2), 'VerticalMid'),\n",
    "                ],\n",
    "                \"Hshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'VerticalMid'),\n",
    "                    ((1, 0), 'VerticalMid'),\n",
    "                    ((1, 2), 'VerticalMid'),\n",
    "                    ((2, 1), 'VerticalMid'),\n",
    "                    ((0, 2), 'Parallel'),\n",
    "                    ((2, 0), 'Parallel'),\n",
    "                ],\n",
    "                \"Rect\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'VerticalEdge'),\n",
    "                    ((1, 0), 'VerticalEdge'),\n",
    "                    ((1, 2), 'VerticalEdge'),\n",
    "                    ((2, 1), 'VerticalEdge'),\n",
    "                    ((2, 3), 'VerticalEdge'),\n",
    "                    ((3, 2), 'VerticalEdge'),\n",
    "                    ((0, 3), 'VerticalEdge'),\n",
    "                    ((3, 0), 'VerticalEdge'),\n",
    "                    ((0, 2), 'Parallel'),\n",
    "                    ((2, 0), 'Parallel'),\n",
    "                    ((1, 3), 'Parallel'),\n",
    "                    ((3, 1), 'Parallel'),\n",
    "                ],\n",
    "                \"RectE1a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                    ((1, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 1), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE1b\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                    ((1, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 1), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE2a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE2b\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((0, 2), \"IsNonOverlapXY\"),\n",
    "                    ((2, 0), \"IsNonOverlapXY\"),\n",
    "                ],\n",
    "                \"RectE3a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((2, 0), \"IsInside\"),\n",
    "                    ((0, 2), \"IsEnclosed\"),\n",
    "                ],\n",
    "                \"RectE3b\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0, 1), \"IsInside\"),\n",
    "                    ((1, 0), \"IsEnclosed\"),\n",
    "                    ((2, 1), \"IsInside\"),\n",
    "                    ((1, 2), \"IsEnclosed\"),\n",
    "                    ((2, 0), \"IsInside\"),\n",
    "                    ((0, 2), \"IsEnclosed\"),\n",
    "                ],\n",
    "                \"RectEconcept\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                ],\n",
    "                \"Graph1\": [\n",
    "                    (0, \"Red\"),\n",
    "                    ((0,1), \"SameColor\"),\n",
    "                    ((1,2), \"SameShape\"),\n",
    "                    ((1,0), \"SameColor\"),\n",
    "                    ((2,1), \"SameShape\"),\n",
    "                ],\n",
    "                \"Graph2\": [\n",
    "                    (0, \"Large\"),\n",
    "                    ((0,1), \"SameSize\"),\n",
    "                    ((0,2), \"SameColor\"),\n",
    "                    ((1,0), \"SameSize\"),\n",
    "                    ((2,0), \"SameColor\"),   \n",
    "                ],\n",
    "                \"Graph3\": [\n",
    "                    (0, \"Cube\"),\n",
    "                    ((0,1), \"SameShape\"),\n",
    "                    ((1,2), \"SameSize\"),\n",
    "                    ((1,0), \"SameShape\"),\n",
    "                    ((2,1), \"SameSize\"),\n",
    "                ],\n",
    "            }\n",
    "    else:\n",
    "        if not is_bidirectional_re:\n",
    "            graph_gt_dict = {\n",
    "                \"Eshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                    ((0, 2), 'Vertical'),\n",
    "                    ((0, 3), 'Vertical'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((1, 3), 'Parallel'),\n",
    "                    ((2, 3), 'Parallel'),\n",
    "                ],\n",
    "                \"Fshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                    ((0, 2), 'Vertical'),\n",
    "                    ((1, 2), 'Parallel'),  \n",
    "                ],\n",
    "                \"Lshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                ],\n",
    "                \"Tshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                ],\n",
    "                \"Cshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                    ((0, 2), 'Vertical'),\n",
    "                    ((1, 2), 'Parallel'),  \n",
    "                ],\n",
    "                \"Ashape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                    ((0, 2), 'Vertical'),\n",
    "                    ((0, 3), 'Parallel'),\n",
    "                    ((1, 2), 'Parallel'),\n",
    "                    ((1, 3), 'Vertical'),\n",
    "                    ((2, 3), 'Vertical'),\n",
    "                ],\n",
    "                \"Hshape\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                    ((1, 2), 'Vertical'),\n",
    "                    ((0, 2), 'Parallel'),\n",
    "                ],\n",
    "                \"Rect\": [\n",
    "                    (0, \"Line\"),\n",
    "                    (1, \"Line\"),\n",
    "                    (2, \"Line\"),\n",
    "                    (3, \"Line\"),\n",
    "                    ((0, 1), 'Vertical'),\n",
    "                    ((1, 2), 'Vertical'),\n",
    "                    ((2, 3), 'Vertical'),\n",
    "                    ((0, 3), 'Vertical'),\n",
    "                    ((0, 2), 'Parallel'),\n",
    "                    ((1, 3), 'Parallel'),\n",
    "                ],\n",
    "                \"RectE1a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0,1), \"IsInside\"),\n",
    "                ],\n",
    "                \"RectE2a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0,1), \"IsInside\"),\n",
    "                    ((2,1), \"IsInside\"),\n",
    "                ],\n",
    "                \"RectE3a\": [\n",
    "                    (0, \"Rect\"),\n",
    "                    (1, \"Rect\"),\n",
    "                    (2, \"Eshape\"),\n",
    "                    ((0,1), \"IsInside\"),\n",
    "                    ((2,1), \"IsInside\"),\n",
    "                    ((2,0), \"IsInside\"),\n",
    "                ],\n",
    "            }\n",
    "        else:\n",
    "            raise\n",
    "    graph = graph_gt_dict[c_type]\n",
    "    if not is_concept:\n",
    "        graph = [ele for ele in graph if not isinstance(ele[0], Number)]\n",
    "    if not is_relation:\n",
    "        graph = [ele for ele in graph if not isinstance(ele[0], tuple)]\n",
    "    return graph\n",
    "\n",
    "\n",
    "def get_reverse_re(r_type):\n",
    "    \"\"\"Get the relation in the reverse direction.\"\"\"\n",
    "    Dict = {\n",
    "        \"Parallel\": \"Parallel\",\n",
    "        \"Vertical\": \"Vertical\",\n",
    "        \"VerticalMid\": \"VerticalMid\",\n",
    "        \"VerticalEdge\": \"VerticalEdge\",\n",
    "        \"SameColor\": \"SameColor\",\n",
    "        \"SameShape\": \"SameShape\",\n",
    "        \"SameAll\": \"SameAll\",\n",
    "        \"IsInside\": \"IsOutside\",\n",
    "        \"SameRow\": \"SameRow\",\n",
    "        \"SameCol\": \"SameCol\",\n",
    "        \"IsTouch\": \"IsTouch\",\n",
    "    }\n",
    "    return Dict[r_type]\n",
    "\n",
    "\n",
    "def get_penalty_dict(graph, threshold=0.35):\n",
    "    penalty_dict = {}\n",
    "    penalty_value_dict = {}\n",
    "    invalid_relations = []\n",
    "    for item in graph:\n",
    "        if isinstance(item[0], tuple):\n",
    "            if item[2] > threshold:\n",
    "                invalid_relations.append(item)\n",
    "                record_data(penalty_dict, [1, 1], [item[0][0], item[0][1]])\n",
    "                record_data(penalty_value_dict, [item[2], item[2]], [item[0][0], item[0][1]])\n",
    "    penalty_dict = transform_dict(penalty_dict, \"sum\")\n",
    "    penalty_value_dict = transform_dict(penalty_value_dict, \"sum\")\n",
    "    return penalty_dict, penalty_value_dict\n",
    "\n",
    "\n",
    "def remove_node(graph, key):\n",
    "    id_to_remove = []\n",
    "    removed_items = []\n",
    "    for id, item in enumerate(graph):\n",
    "        if isinstance(item[0], Number):\n",
    "            if item[0] == key:\n",
    "                id_to_remove.append(id)\n",
    "        else:\n",
    "            assert isinstance(item, tuple)\n",
    "            if item[0][0] == key or item[0][1] == key:\n",
    "                id_to_remove.append(id)\n",
    "    for id in reversed(id_to_remove):\n",
    "        removed_items.insert(0, graph.pop(id))\n",
    "    return graph, removed_items\n",
    "\n",
    "\n",
    "def filter_graph_with_threshold(graph, mode=\"greedy\", threshold=0.35):\n",
    "    \"\"\"Filter graph based on threshold.\"\"\"\n",
    "    def get_argmax_key(Dict):\n",
    "        id = np.argmax(list(Dict.values()))\n",
    "        key = list(Dict)[id]\n",
    "        return key\n",
    "    graph = deepcopy(graph)\n",
    "    if threshold is None:\n",
    "        return graph, []\n",
    "    removed_items = []\n",
    "    if mode == \"greedy\":\n",
    "        for i in range(len(graph)):\n",
    "            penalty_dict, penalty_value_dict = get_penalty_dict(graph, threshold=threshold)\n",
    "            if len(penalty_value_dict) == 0:\n",
    "                break\n",
    "            key_to_remove = get_argmax_key(penalty_value_dict)\n",
    "            graph, removed_item = remove_node(graph, key_to_remove)\n",
    "            removed_items += removed_item\n",
    "    else:\n",
    "        raise\n",
    "    return graph, removed_items\n",
    "\n",
    "\n",
    "def filter_graph_with_masks(graph, masks_is_invalid):\n",
    "    \"\"\"Filter the graph with given masks_is_invalid.\n",
    "\n",
    "    Args:\n",
    "        masks_is_invalid: has the format of e.g. [True, False, True, False, False].\n",
    "    \"\"\"\n",
    "    invalid_ids = np.where(masks_is_invalid)[0]\n",
    "    removed_indices = []\n",
    "    for i in range(len(graph)):\n",
    "        node_or_edge = graph[i][0]\n",
    "        if isinstance(node_or_edge, tuple):\n",
    "            # An edge:\n",
    "            to_remove = False\n",
    "            for k in invalid_ids:\n",
    "                if k in node_or_edge:\n",
    "                    to_remove = True\n",
    "                    break\n",
    "        else:\n",
    "            # A node:\n",
    "            to_remove = node_or_edge in invalid_ids\n",
    "        if to_remove:\n",
    "            removed_indices.append(i)\n",
    "    new_graph = [graph[i] for i in range(len(graph)) if i not in removed_indices]\n",
    "    return new_graph\n",
    "\n",
    "\n",
    "def random_mask_obj(input, prob):\n",
    "    p = 1 / 2 / (1 - prob)\n",
    "    mask = input.argmax(1) != 0\n",
    "    out = mask.float() * (torch.rand(mask.shape) * p).round().to(mask.device)\n",
    "    return out.round()[None]\n",
    "\n",
    "\n",
    "def random_obj(input, concept_prob_dict):\n",
    "    \"\"\"Choose random object based on the concept type.\"\"\"\n",
    "    list_of_objs = find_connected_components_colordiff(onehot_to_RGB(input).squeeze(), is_mask=True)\n",
    "    Dict = {}\n",
    "    for i, obj in enumerate(list_of_objs):\n",
    "        mask = (obj[0].max(0)[0] != 0).float()\n",
    "        c_type = classify_concept(mask)\n",
    "        Dict[i] = {\"c_type\": c_type, \"prob\": concept_prob_dict[c_type]}\n",
    "    prob_list = np.array([ele[\"prob\"] for ele in Dict.values()])\n",
    "    prob_list = prob_list / prob_list.sum()\n",
    "    id = np.random.choice(len(prob_list), p=prob_list)\n",
    "    mask = torch.FloatTensor(list_of_objs[id][2]).to(device)\n",
    "    return mask[None,None]\n",
    "\n",
    "\n",
    "def get_rand_graph(freq_dict, is_new_vertical=False):\n",
    "    line_num, counts = np.unique(freq_dict[\"Line\"], return_counts=True)\n",
    "    line_freq = counts / counts.sum()\n",
    "    lines = np.random.choice(line_num, p=line_freq)\n",
    "    graph_list = []\n",
    "    for i in range(lines):\n",
    "        graph_list.append((i, \"Line\"))\n",
    "    id_pairs = list(zip(*get_triu_ids(lines)))\n",
    "    for id_pair in id_pairs:\n",
    "        # re_chosen = np.random.choice([\"VerticalMid\", \"VerticalEdge\", \"Parallel\"], p=[1/3, 1/3, 1/3])\n",
    "        if is_new_vertical:\n",
    "            re_chosen = np.random.choice([\"VerticalMid\", \"VerticalEdge\", \"Parallel\"], p=freq_ver_para)\n",
    "        else:\n",
    "            re_chosen = np.random.choice([\"Vertical\", \"Parallel\"], p=freq_ver_para)\n",
    "        graph_list.append((id_pair, re_chosen))\n",
    "    return graph_list\n",
    "\n",
    "\n",
    "def get_c_type_from_data(data):\n",
    "    if isinstance(data[2], tuple):\n",
    "        assert len(np.unique(data[2])) == 1\n",
    "        # Requiring that all data has the same positive labels:\n",
    "        assert data[2][0] == data[3][0][\"obj_spec\"][0][0][1].split(\"_\")[0]\n",
    "        return data[2][0]\n",
    "    else:\n",
    "        assert data[2] == data[3][\"obj_spec\"][0][0][1].split(\"_\")[0]\n",
    "        return data[2]\n",
    "\n",
    "\n",
    "def random_obj(input, concept_prob_dict=None):\n",
    "    \"\"\"Choose random object based on the concept type.\"\"\"\n",
    "    list_of_objs = find_connected_components_colordiff(onehot_to_RGB(input).squeeze(), is_mask=True)\n",
    "    id = np.random.choice(len(list_of_objs))\n",
    "    mask = torch.FloatTensor(list_of_objs[id][2]).to(device)\n",
    "    return mask[None,None]\n",
    "\n",
    "\n",
    "def shorten_graph(graph):\n",
    "    return [ele[:3] for ele in graph]\n",
    "\n",
    "\n",
    "def is_mask_invalid(mask, threshold_pixels=0):\n",
    "    n_pixels = (mask.round() > 0).long().sum().item()\n",
    "    return n_pixels <= threshold_pixels\n",
    "\n",
    "\n",
    "def get_fewshot_gt_idx(dataloader):\n",
    "    gt_idx_list = []\n",
    "    for data in dataloader:\n",
    "        concept_id = data[2][\"concept_id\"][0]\n",
    "        example_concept_ids = [ele[0] for ele in data[1][2]]\n",
    "        gt_idx = example_concept_ids.index(concept_id)\n",
    "        gt_idx_list.append(gt_idx)\n",
    "    gt_idx_list = np.array(gt_idx_list)\n",
    "    return gt_idx_list\n",
    "\n",
    "\n",
    "def update_default_hyperparam(Dict):\n",
    "    \"\"\"Default hyperparameters for previous experiments, after adding these new options.\"\"\"\n",
    "    default_param = {\n",
    "        \"is_round_mask\": True,\n",
    "        \"SGLD_is_penalize_lower\": \"False\",\n",
    "        \"val_batch_size\": 1,\n",
    "    }\n",
    "    for key, item in default_param.items():\n",
    "        if key not in Dict:\n",
    "            Dict[key] = item\n",
    "    return Dict\n",
    "\n",
    "\n",
    "def update_default_hyperparam_generalization(Dict):\n",
    "    \"\"\"Default hyperparameters for previous experiments, after adding these new options.\"\"\"\n",
    "    default_param = {\n",
    "        \"init\": \"random\",\n",
    "        \"SGLD_is_penalize_lower_seq\": \"None\",\n",
    "        \"lambd\": 0.005,\n",
    "        \"infer_order\": \"simul\",\n",
    "        \"is_bidirectional_re\": False,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"max_n_distractors\": 3,\n",
    "        \"allow_connect\": True,\n",
    "        \"is_concept\": True,\n",
    "        \"is_relation\": True,\n",
    "        \"is_proper_size\": False,\n",
    "        \"is_harder_distractor\": False,\n",
    "        \"id\": \"None\",\n",
    "    }\n",
    "    for key, item in default_param.items():\n",
    "        if key not in Dict:\n",
    "            Dict[key] = item\n",
    "    return Dict\n",
    "\n",
    "\n",
    "def get_selector_SGLD(\n",
    "    selector,\n",
    "    input,\n",
    "    args,\n",
    "    neg_mask=None,\n",
    "    mask_exclude=None,\n",
    "    ebm_target=\"mask\",\n",
    "    ensemble_size=16,\n",
    "    sample_step=150,\n",
    "    topk=1,\n",
    "    isplot=1,\n",
    "    init=\"mask-random\",\n",
    "    **kwargs\n",
    "):\n",
    "    \"\"\"\n",
    "    Performs SGLD on the mask or img, given the selector.\n",
    "\n",
    "    Returns:\n",
    "        imgs_top: each element has shape of [topk, B, 10, H, W]\n",
    "        masks_top: each element has shape of [topk, B, 1, H, W]\n",
    "    \"\"\"\n",
    "    def init_neg_mask(pos_img, init, args):\n",
    "        \"\"\"Initialize negative mask\"\"\"\n",
    "        assert not isinstance(pos_img, tuple) and len(pos_img.shape) == 4\n",
    "        if init == \"mask-random\":\n",
    "            assert pos_img.shape[1] == 10\n",
    "            pos_img_l = repeat_n(pos_img, n_repeats=ensemble_size)\n",
    "            neg_mask = tuple(torch.rand(pos_img_l.shape[0], 1, *pos_img_l.shape[2:]).to(pos_img.device) * (pos_img_l.argmax(1)[:, None] != 0) for k in range(selector.mask_arity))\n",
    "            for k in range(len(neg_mask)):\n",
    "                neg_mask[k].requires_grad = True\n",
    "        elif init == \"random\":\n",
    "            neg_mask = None\n",
    "        else:\n",
    "            raise\n",
    "        return neg_mask\n",
    "\n",
    "    if neg_mask is None:\n",
    "        neg_mask = init_neg_mask(input, init, args)\n",
    "    is_grad = False\n",
    "    args = deepcopy(args)\n",
    "    \n",
    "    for key, value in kwargs.items():\n",
    "        if value is not None:\n",
    "            setattr(args, key, value)\n",
    "    args.ebm_target = ebm_target\n",
    "\n",
    "    (img_ensemble, neg_mask_ensemble, _, _, _), neg_out_list_ensemble, info_ensemble = neg_mask_sgd_ensemble(\n",
    "        selector, input, neg_mask, c_repr=None, z=None, zgnn=None, wtarget=None, args=args,\n",
    "        mask_info={\"mask_exclude\": mask_exclude} if mask_exclude is not None else None,\n",
    "        ensemble_size=ensemble_size, is_grad=is_grad,\n",
    "        out_mode=\"all-sorted\",\n",
    "        return_history=True,\n",
    "        record_interval=2,\n",
    "    )\n",
    "    masks_top = tuple(neg_mask_ensemble[k][:topk] for k in range(len(neg_mask_ensemble)))\n",
    "    if ebm_target in [\"image+mask\"]:\n",
    "        imgs_top = img_ensemble[:topk]\n",
    "    else:\n",
    "        imgs_top = None\n",
    "    energy_mask = {}\n",
    "    info = {}\n",
    "    energy = neg_out_list_ensemble[-1][:topk]\n",
    "    for key in selector.info:\n",
    "        energy_mask[key] = selector.info[key][:topk]  # each with [topk, B]\n",
    "    # energy_mask = np.stack(list(energy_mask.values()), -1)\n",
    "    info[\"mutual_exclusive_list\"] = info_ensemble[\"mutual_exclusive_list\"][-1][:topk]\n",
    "    info[\"energy\"] = energy\n",
    "    info[\"energy_mask\"] = energy_mask\n",
    "    info[\"mask_list\"] = tuple(mask_list[:,:topk] for mask_list in info_ensemble[\"neg_mask_list\"])\n",
    "    refer_node_names = selector.refer_node_names if selector.refer_node_names is not None else list(selector.nodes)\n",
    "    if isplot >= 1:\n",
    "        # energy_str = [\"{:.4f}\".format(ele) for ele in energy]\n",
    "        if ebm_target == \"mask\":\n",
    "            for j in range(len(input)):\n",
    "                if input[j].shape[-3] == 10:\n",
    "                    visualize_matrices([input[j].argmax(0)], use_color_dict=True)\n",
    "                elif input[j].shape[-3] == 3:\n",
    "                    visualize_matrices([input[j]], use_color_dict=False)\n",
    "                else:\n",
    "                    raise\n",
    "                energy_mask_j = {key: value[:,j] for key, value in energy_mask.items()}\n",
    "                pp.pprint(energy_mask_j)\n",
    "                # print([\"[E={:.4f}]\".format(ele) if list(selector.topological_sort)[k] in refer_node_names else \"E={:.4f}\".format(ele) for k, ele in enumerate(energy_mask[j])])\n",
    "                if \"mutual_exclusive_list\" in info:\n",
    "                    print(\"mutual_exclusive: {}\".format(info[\"mutual_exclusive_list\"][:,j]))\n",
    "                for k in range(topk):\n",
    "                    print(\"top {}, energy: {:.3f}   mutual_exclusive: {:.3f}:\".format(k, energy[k,j], info[\"mutual_exclusive_list\"][k,j]))\n",
    "                    subtitles = [\"E={:.3f}\".format(value[k]) for key, value in energy_mask_j.items() if \"(\" not in key]\n",
    "                    if len(subtitles) < len(masks_top):\n",
    "                        subtitles += [None for _ in range(len(masks_top) - len(subtitles))]\n",
    "                    plot_matrices(torch.stack(masks_top)[:,k,j].squeeze(-3), scale_limit=(0,1), images_per_row=6,\n",
    "                                  subtitles=subtitles)\n",
    "        elif ebm_target == \"image+mask\":\n",
    "            for j in range(len(input)):\n",
    "                energy_mask_j = {key: value[:,j] for key, value in energy_mask.items()}\n",
    "                pp.pprint(energy_mask_j)\n",
    "                if \"mutual_exclusive_list\" in info:\n",
    "                    print(\"mutual_exclusive: {}\".format(info[\"mutual_exclusive_list\"][:,j]))\n",
    "                for k in range(topk):\n",
    "                    visualize_matrices([imgs_top[k,j].argmax(0)], is_color_dict=True, subtitles=[\"E={:.4f}\".format(energy[k,j])])\n",
    "                    plot_matrices(torch.stack(masks_top)[:,k,j].squeeze(-3), scale_limit=(0,1), images_per_row=6,\n",
    "                                  subtitles=[\"E={:.3f}\".format(value[k]) for key, value in energy_mask_j.items()])\n",
    "        else:\n",
    "            raise\n",
    "    return imgs_top, masks_top, info\n",
    "\n",
    "\n",
    "def parse_selector_from_image(\n",
    "    input,\n",
    "    args,\n",
    "    keys_dict,\n",
    "    topk=3,\n",
    "    isplot=False,\n",
    "    infer_order=\"simul\",\n",
    "    init=\"mask-random\",\n",
    "    threshold_pixels=0,\n",
    "):\n",
    "    \"\"\"Obtain selector from parsing the image.\"\"\"\n",
    "    if infer_order == \"simul\":\n",
    "        selector = get_selector_for_parsing(keys_dict, ebm_dict, CONCEPTS, OPERATORS)\n",
    "        # Parse the graph\n",
    "        _, masks_top, info = get_selector_SGLD(\n",
    "            selector, input, args,\n",
    "            ebm_target=\"mask\",\n",
    "            init=init,\n",
    "            SGLD_object_exceed_coef=0,\n",
    "            SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "            SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "            ensemble_size=args.ensemble_size,\n",
    "            sample_step=args.sample_step,\n",
    "            topk=topk,\n",
    "            isplot=isplot,\n",
    "        )\n",
    "    elif infer_order == \"seq\":\n",
    "        # Parse concepts sequentially:\n",
    "        masks_top_sum = None\n",
    "        if input.shape[1] == 10:\n",
    "            assert len(input.shape) == 4\n",
    "            masks_all_gt = input[:,:1] != 1\n",
    "        masks_collect = []\n",
    "        key_list = []\n",
    "        for key in keys_dict:\n",
    "            key_list += [key] * keys_dict[key]\n",
    "        key_list = np.random.permutation(key_list)\n",
    "        for key in key_list:\n",
    "            selector = get_selector_for_parsing({key: 1}, ebm_dict, CONCEPTS, OPERATORS)\n",
    "            _, masks_top, info = get_selector_SGLD(\n",
    "                selector, input, args,\n",
    "                ebm_target=\"mask\",\n",
    "                mask_exclude=masks_top_sum,\n",
    "                init=init,\n",
    "                SGLD_object_exceed_coef=0,\n",
    "                SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "                SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "                SGLD_is_penalize_lower=args.SGLD_is_penalize_lower_seq,\n",
    "                ensemble_size=args.ensemble_size,\n",
    "                sample_step=args.sample_step,\n",
    "                topk=1,\n",
    "                isplot=isplot,\n",
    "            )\n",
    "            for mask in masks_top:\n",
    "                if mask.round().sum() > 0:\n",
    "                    masks_collect.append(repeat_n(mask.squeeze(1), n_repeats=args.ensemble_size))\n",
    "            if masks_top_sum is None:\n",
    "                masks_top_sum = torch.stack([mask.detach() for mask in masks_top]).sum(0).clamp(0, 1).round().squeeze(1)\n",
    "            else:\n",
    "                masks_top_sum = (masks_top_sum + torch.stack([mask.detach() for mask in masks_top]).sum(0).squeeze(1).round()).clamp(0, 1)\n",
    "            if isplot:\n",
    "                print(\"masks_sum:\")\n",
    "                plot_matrices(masks_top_sum.squeeze(1))\n",
    "            if input.shape[1] == 10:\n",
    "                if np.clip(to_np_array(masks_all_gt) - to_np_array(masks_top_sum), a_min=0, a_max=None).sum() == 0:\n",
    "                    if isplot:\n",
    "                        print(\"All objects accounted for. Break.\")\n",
    "                        break\n",
    "        # Fine-tune with all concepts:\n",
    "        selector = get_selector_for_parsing(keys_dict, ebm_dict, CONCEPTS, OPERATORS)\n",
    "        # Parse the graph:\n",
    "        _, masks_top, info = get_selector_SGLD(\n",
    "            selector, input, args,\n",
    "            neg_mask=tuple(masks_collect),\n",
    "            ebm_target=\"mask\",\n",
    "            init=init,\n",
    "            lambd_start=args.lambd,\n",
    "            lambd=args.lambd,\n",
    "            SGLD_object_exceed_coef=0,\n",
    "            SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "            SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "            ensemble_size=args.ensemble_size,\n",
    "            sample_step=min(60, args.sample_step),\n",
    "            topk=topk,\n",
    "            isplot=isplot,\n",
    "        )\n",
    "\n",
    "    else:\n",
    "        raise Exception(\"infer_order '{}' is not valid!\".format(infer_order))\n",
    "    graph_dict = {k: [] for k in range(topk)}\n",
    "    graph_trimmed_dict = {}\n",
    "    input_expand = input.expand(topk, *input.shape[1:])\n",
    "    for i, node_name in enumerate(selector.topological_sort):\n",
    "        c_type = node_name.split(\":\")[-1]\n",
    "        E_c = ebm_dict[c_type](input_expand, (masks_top[i].squeeze(1),))\n",
    "        for k in range(topk):\n",
    "            graph_dict[k].append((i, c_type, E_c[k].item()))\n",
    "\n",
    "    id_pairs = list(zip(*get_triu_ids(len(masks_top))))\n",
    "    for id0, id1 in id_pairs:\n",
    "        E_re = {}\n",
    "        for re_key in re_keys:\n",
    "            if args.is_round_mask:\n",
    "                # Each value has shape of [topk, 1]\n",
    "                if args.is_bidirectional_re:\n",
    "                    E_re[re_key] = to_np_array(ebm_dict[re_key]((input_expand, input_expand), (masks_top[id0].round().squeeze(1), masks_top[id1].round().squeeze(1))) +\n",
    "                                               ebm_dict[get_reverse_re(re_key)]((input_expand, input_expand), (masks_top[id1].round().squeeze(1), masks_top[id0].round().squeeze(1))))\n",
    "                else:\n",
    "                    E_re[re_key] = to_np_array(ebm_dict[re_key]((input_expand, input_expand), (masks_top[id0].round().squeeze(1), masks_top[id1].round().squeeze(1))))\n",
    "            else:\n",
    "                if args.is_bidirectional_re:\n",
    "                    E_re[re_key] = to_np_array(ebm_dict[re_key]((input_expand, input_expand), (masks_top[id0].squeeze(1), masks_top[id1].squeeze(1))) +\n",
    "                                               ebm_dict[get_reverse_re(re_key)]((input_expand, input_expand), (masks_top[id1].squeeze(1), masks_top[id0].squeeze(1))))\n",
    "                else:\n",
    "                    E_re[re_key] = to_np_array(ebm_dict[re_key]((input_expand, input_expand), (masks_top[id0].squeeze(1), masks_top[id1].squeeze(1))))\n",
    "        E_re_all = np.concatenate(list(E_re.values()), -1)\n",
    "        idx_argmin = E_re_all.argmin(-1)\n",
    "        re_argmin = np.array(list(E_re))[idx_argmin].tolist()\n",
    "        E_min = np.take_along_axis(E_re_all, indices=idx_argmin[:,None], axis=-1).squeeze()\n",
    "        for k in range(topk):\n",
    "            graph_dict[k].append(((id0, id1), re_argmin[k], E_min[k]))\n",
    "\n",
    "    for k in range(topk):\n",
    "        masks_is_invalid = [is_mask_invalid(mask[k], threshold_pixels=threshold_pixels) for mask in masks_top]\n",
    "        graph_trimmed_dict[k] = filter_graph_with_masks(graph_dict[k], masks_is_invalid=masks_is_invalid)\n",
    "    return graph_trimmed_dict, graph_dict, tuple(mask.squeeze(1) for mask in masks_top), info"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6456298c-78b2-4304-99b0-cf74eec19969",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make sure that the initialization of CONCEPTS and OPERATORS (including their embedding) is after setting the seed\n",
    "set_seed(1)\n",
    "from zeroc.concepts_shapes import OPERATORS, CONCEPTS, load_task, seperate_concept\n",
    "try:\n",
    "    get_ipython().run_line_magic('matplotlib', 'inline')\n",
    "    is_jupyter = True\n",
    "except:\n",
    "    is_jupyter = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ab0504f-dbd0-4c1c-b3b4-bfbf7613252e",
   "metadata": {},
   "source": [
    "# 2. Inference:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03e83a7f-ed59-4bf9-b6ce-51b4df0c5dbd",
   "metadata": {},
   "source": [
    "### 2.1 Loading atomic EBMs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d43ed5f4-8b03-4e8d-94ad-377976118abf",
   "metadata": {},
   "outputs": [],
   "source": [
    "is_load = False\n",
    "date_time = f\"{datetime.now().month}-{datetime.now().day}\"\n",
    "if is_load:\n",
    "    all_dict = pload(EXP_PATH + \"/generalization/all_dict_April11.p\")\n",
    "    models = all_dict[\"models\"]\n",
    "    CONCEPTS = all_dict[\"CONCEPTS\"]\n",
    "    OPERATORS = all_dict[\"OPERATORS\"]\n",
    "    SGLD_args = all_dict[\"args\"]\n",
    "else:\n",
    "    if \"models\" not in locals():\n",
    "        models = {}\n",
    "    device = \"cpu\"\n",
    "\n",
    "    all_dict = {\n",
    "        \"models\": models,\n",
    "        \"CONCEPTS\": CONCEPTS,\n",
    "        \"OPERATORS\": OPERATORS,\n",
    "        # \"args\": info[\"args\"],\n",
    "    }\n",
    "    for key, model in models.items():\n",
    "        model.to(\"cpu\")\n",
    "    make_dir(EXP_PATH + f\"/generalization/all_dict_{date_time}.p\")\n",
    "    pdump(all_dict, EXP_PATH + f\"/generalization/all_dict_{date_time}.p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08fb0bf2-e772-4ae1-8b0c-2a26258549e3",
   "metadata": {},
   "source": [
    "### 2.2 Settings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "011b001f-14bc-45c0-abf4-fb0c9b5ad6dc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description='PDE argparse.')\n",
    "parser.add_argument('--evaluation_type', type=str, help='Evaluation type')\n",
    "parser.add_argument('--dataset', type=str, default=\"h-r^2ai+2a+3ai+3a+3b:SameShape+SameColor(Line+Rect+RectSolid+Lshape)\", help='dataset')\n",
    "parser.add_argument('--canvas_size', type=int, default=16, help='Canvas_size')\n",
    "parser.add_argument('--canvas_size_3D', type=int, default=32, help='Canvas_size')\n",
    "parser.add_argument('--color_avail', type=str, default=\"1,2\", help='color_avail.')\n",
    "parser.add_argument('--max_n_distractors', type=int, default=3, help='Maximum number of distractors for grounding.')\n",
    "parser.add_argument('--min_n_distractors', type=int, default=0, help='Minimum number of distractors for grounding.')\n",
    "parser.add_argument('--allow_connect', type=str2bool, nargs='?', const=True, default=True,\n",
    "                        help='Whether or not to allow objects to connect in the image.')\n",
    "parser.add_argument('--is_proper_size', type=str2bool, nargs='?', const=True, default=False,\n",
    "                        help='If True, the main character for grounding will have a proper size.')\n",
    "parser.add_argument('--is_harder_distractor', type=str2bool, nargs='?', const=True, default=False,\n",
    "                        help='If True, will use harder distractors.')\n",
    "parser.add_argument('--model_type', type=str, default=\"hc-ebm\", help='Model type.')\n",
    "parser.add_argument('--SGLD_mutual_exclusive_coef', type=float, default=0, help='SGLD_mutual_exclusive_coef.')\n",
    "parser.add_argument('--SGLD_pixel_entropy_coef', type=float, default=0, help='SGLD_pixel_entropy_coef.')\n",
    "parser.add_argument('--SGLD_is_anneal', type=str2bool, nargs='?', const=True, default=True,\n",
    "                        help='If True, will anneal the SGLD coefficients.')\n",
    "parser.add_argument('--SGLD_is_penalize_lower', type=str, default=\"False\", help='if True or \"True\", will penalize that the sum is less than 1. If \"False\" or False, will not. If \"obj\", will only penalize on the object locations (if n_channels==10)..')\n",
    "parser.add_argument('--SGLD_is_penalize_lower_seq', type=str, default=\"None\", help='if True or \"True\", will penalize that the sum is less than 1. If \"False\" or False, will not. If \"obj\", will only penalize on the object locations (if n_channels==10)..')\n",
    "parser.add_argument('--concept_model_hash', type=str, default=\"None\", help='hash key for the concept model.')\n",
    "parser.add_argument('--relation_model_hash', type=str, default=\"None\", help='hash key for the relation model.')\n",
    "parser.add_argument('--concept_model_hash_3D', type=str, default=\"None\", help='hash key for the concept model.')\n",
    "parser.add_argument('--relation_model_hash_3D', type=str, default=\"None\", help='hash key for the relation model.')\n",
    "parser.add_argument('--concept_load_id', type=str, default=\"best\", help='\"best\" or a number string.')\n",
    "parser.add_argument('--relation_load_id', type=str, default=\"best\", help='\"best\" or a number string.')\n",
    "parser.add_argument('--concept_load_id_3D', type=str, default=\"best\", help='\"best\" or a number string.')\n",
    "parser.add_argument('--relation_load_id_3D', type=str, default=\"best\", help='\"best\" or a number string.')\n",
    "parser.add_argument('--is_new_vertical', type=str2bool, nargs='?', const=True, default=True,\n",
    "                        help='If True, will use VertMid, VertEdge.')\n",
    "parser.add_argument('--is_concept', type=str2bool, nargs='?', const=True, default=True,\n",
    "                        help='If False, will remove the concept EBMs in the selector.')\n",
    "parser.add_argument('--is_relation', type=str2bool, nargs='?', const=True, default=True,\n",
    "                        help='If False, will remove the relation EBMs in the selector.')\n",
    "parser.add_argument('--is_bidirectional_re', type=str2bool, nargs='?', const=True, default=False,\n",
    "                        help='If True, each relation will be evaluated on both directions.')\n",
    "parser.add_argument('--is_round_mask', type=str2bool, nargs='?', const=True, default=True,\n",
    "                        help='If True, will use VertMid, VertEdge.')\n",
    "parser.add_argument('--val_batch_size', type=int, default=1, \n",
    "                        help='batch_size for validation.')\n",
    "parser.add_argument('--val_n_examples', type=int, default=400, \n",
    "                        help='batch_size for validation.')\n",
    "parser.add_argument('--is_analysis', type=str2bool, nargs='?', const=True, default=False,\n",
    "                        help='If True, will perform analysis.')\n",
    "parser.add_argument('--ensemble_size', type=int, default=16, help='ensemble_size')\n",
    "parser.add_argument('--sample_step', type=int, default=150, help='ensemble_size')\n",
    "parser.add_argument('--lambd', type=float, default=0.005, help='noise scale at the end.')\n",
    "parser.add_argument('--infer_order', type=str, default=\"simul\", help='Choose from \"parallel\", \"sequential\".')\n",
    "parser.add_argument('--init', type=str, default=\"random\", help='Choose from \"random\", \"mask-random\"')\n",
    "parser.add_argument('--load_parse_src', type=str, default=\"gt\", help='Choose from \"gt\", or a path to a parse file.')\n",
    "parser.add_argument('--inspect_interval', type=int, default=20, help='inspect interval')\n",
    "parser.add_argument('--seed', type=int, default=2, help='color_avail.')\n",
    "parser.add_argument('--topk', type=int, default=16, help='color_avail.')\n",
    "parser.add_argument('--id', type=str, default=\"0\", help='Id')\n",
    "parser.add_argument('--gpuid', type=str, help='Id')\n",
    "parser.add_argument('--date_time', type=str, help='Month and date')\n",
    "parser.add_argument('--data_range', type=str, default=\"None\", help='\"None\" or \"100:200\"')\n",
    "parser.add_argument('--verbose', type=int, default=1, help='Verbose.')\n",
    "\n",
    "try:\n",
    "    get_ipython().run_line_magic('matplotlib', 'inline')\n",
    "    args = parser.parse_args([])\n",
    "    args.min_n_distractors = 1\n",
    "    args.max_n_distractors = 2\n",
    "    args.is_proper_size = True\n",
    "    # args.dataset = \"pc-Cshape+Eshape+Fshape+Ashape+Hshape+Rect\"\n",
    "    # args.dataset = \"pc-RectE1a+RectE2a+RectE3a\"\n",
    "    # args.dataset = \"pg-Lshape+Tshape+Cshape+Eshape+Fshape+Ashape+Rect^Eshape+Fshape+Ashape+Rect\"\n",
    "    # args.dataset = \"yc-Eshape[5,9]+Fshape[5,9]+Ashape[5,9]\"\n",
    "    args.dataset = \"c-Lshape+Tshape+Cshape+Eshape+Fshape+Ashape\"\n",
    "    # args.dataset = \"c-Eshape+Fshape+Ashape\"\n",
    "    # args.dataset = \"c-RectE1a+RectE2a+RectE3a\"\n",
    "    # args.dataset = \"h-r^2ai+2a+3ai+3a+3b:SameShape+SameColor(Line+Rect+RectSolid+Lshape)\"\n",
    "    # args.evaluation_type = \"pc-parse+classify^parse\" # \"grounding-Lshape\"   # \"parse\", \"grounding\", \"classify\", \"query-1c\", \"unify\", \n",
    "    # args.evaluation_type = \"classify\"\n",
    "    # args.evaluation_type = \"yc-parse+classify\"\n",
    "    # args.dataset = \"c-Lshape+Eshape+Fshape+Ashape\"\n",
    "    args.SGLD_mutual_exclusive_coef = 500\n",
    "    args.SGLD_mask_entropy_coef = 0\n",
    "    args.model_type = \"hc-ebm\"\n",
    "    # args.model_type = \"rand-obj\"  # \"rand-mask-obj\" or \"rand-obj\"\n",
    "    # args.model_type = \"rand-graph\"\n",
    "    args.ensemble_size = 64\n",
    "    args.allow_connect = False\n",
    "    args.inspect_interval = 1\n",
    "    args.sample_step = 150\n",
    "    args.is_new_vertical = True\n",
    "    args.is_concept = True\n",
    "    args.is_bidirectional_re = True\n",
    "    args.canvas_size = 16\n",
    "    args.SGLD_is_anneal = True\n",
    "    args.val_n_examples = 200\n",
    "    args.SGLD_is_penalize_lower = \"False\"\n",
    "    args.SGLD_is_penalize_lower_seq = \"False\"\n",
    "\n",
    "    # args.relation_model_hash = \"1IB9rf++\"  # new3\n",
    "    # args.relation_model_hash = \"BiP7IcmF\" # center^stop:0.1 with r-rmbx\n",
    "    # args.relation_model_hash= \"vcLANlzs\" # center^stop:0.1 with r-mbx\n",
    "    ## max_n_distractors=0:\n",
    "    # args.concept_model_hash = \"VRaV6VTb\" # randpatch rmb, or \"1NHyHlRM\" (rmbx)\n",
    "    # args.relation_model_hash = \"cFyJdFYC\" # randpatch rmb, or \"i2um9RI2\" (rmbx)\n",
    "    args.concept_load_id = \"best\"\n",
    "    args.relation_load_id = \"best\"\n",
    "\n",
    "    ## For classification of Eshape:\n",
    "    args.evaluation_type = \"classify\"\n",
    "    args.dataset = \"c-Eshape+Fshape+Ashape\"\n",
    "    args.canvas_size = 16\n",
    "    args.concept_model_hash = \"fRZtzn33\"\n",
    "    args.relation_model_hash = \"Wfxw19nM\"\n",
    "\n",
    "    # ## For grounding Eshape:\n",
    "    # # args.evaluation_type = \"grounding-RectE1a\"\n",
    "    # args.evaluation_type = \"grounding-Eshape\"\n",
    "    # args.dataset = \"c-Lshape+Tshape+Cshape+Eshape+Fshape+Ashape\"\n",
    "    # args.canvas_size = 16\n",
    "    # args.concept_model_hash = \"fRZtzn33\"\n",
    "    # args.relation_model_hash = \"1IB9rf++\"\n",
    "    # args.relation_load_id = 75\n",
    "    # args.val_n_examples = 400\n",
    "\n",
    "    # # ## For 3D dataset:\n",
    "    # args.evaluation_type = \"yc-parse+classify^classify\"\n",
    "    # args.concept_model_hash = \"fRZtzn33\"\n",
    "    # args.relation_model_hash = \"Wfxw19nM\"\n",
    "    # args.concept_model_hash_3D = \"jk4HQfir\"\n",
    "    # args.relation_model_hash_3D = \"x72bDIyX\"\n",
    "    # args.dataset = \"yc-Eshape[5,9]+Fshape[5,9]+Ashape[5,9]\"\n",
    "    # args.canvas_size = 16\n",
    "    # args.canvas_size_3D = 32\n",
    "    # args.load_parse_src = \"gt\"\n",
    "    # args.load_parse_src = \"evaluation_yc-parse+classify^parse_1-21/evaluation_yc-parse+classify^parse_canvas_16_color_1,2_ex_400_min_0_model_hc-ebm_mutu_500.0_ens_64_sas_150_newv_True_batch_1_con_fRZtzn33_re_Wfxw19nM_bi_True_seed_2_id_None_Hash_mU7ILNWm_turing3.p\"\n",
    "\n",
    "    # ## For RectE:\n",
    "    # args.evaluation_type = \"grounding-RectE1a\"\n",
    "    # args.concept_model_hash = \"AaalzcSD\"\n",
    "    # args.relation_model_hash = \"NAzKCenZ\"\n",
    "    # args.dataset = \"c-Lshape+Tshape+Cshape+Eshape+Fshape+Ashape\"\n",
    "    # args.canvas_size = 20\n",
    "    # args.max_n_distractors = 1\n",
    "    # args.min_n_distractors = 1\n",
    "    # args.is_harder_distractor = False\n",
    "    # # args.concept_load_id = 30\n",
    "    # # args.relation_load_id = 10\n",
    "    \n",
    "    # ## For clevr:\n",
    "    # args.evaluation_type = \"classify\"\n",
    "    # args.dataset = \"u-graph-Graph1+Graph2+Graph3\"\n",
    "    # args.canvas_size = 32\n",
    "    # args.concept_model_hash = \"fuXnlqYJ\"\n",
    "    # args.relation_model_hash = \"LMBhNjDn\"\n",
    "\n",
    "    args.is_round_mask = True\n",
    "    # args.lambd = 0.001\n",
    "    args.infer_order = \"simul\"\n",
    "\n",
    "    # args.val_n_examples = 40\n",
    "    args.val_batch_size = 1\n",
    "    args.is_analysis = True\n",
    "    args.init = \"random\"\n",
    "    args.seed = 2\n",
    "    args.verbose = 2\n",
    "    args.gpuid = \"0\"\n",
    "    args.id = \"test\"\n",
    "    args.date_time = \"{}-{}\".format(datetime.now().month, datetime.now().day)\n",
    "except:\n",
    "    args = parser.parse_args()\n",
    "is_input = False\n",
    "assert args.val_n_examples % args.val_batch_size == 0\n",
    "device = get_device(args)\n",
    "dirname = EXP_PATH + \"/evaluation_{}_{}/\".format(args.evaluation_type, args.date_time)\n",
    "\n",
    "# Model:\n",
    "if args.evaluation_type.startswith(\"query\") or args.evaluation_type.startswith(\"unify\"):\n",
    "    assert args.color_avail == \"1,2\"\n",
    "    assert args.dataset.startswith(\"h-\")\n",
    "    if args.canvas_size == 8:\n",
    "        concept_model = models[\"JjloOXSA\"]\n",
    "        relation_model = models[\"j50bI0im\"]\n",
    "    elif args.canvas_size == 16:\n",
    "        concept_model = models[\"l6Gaygnh\"]\n",
    "        relation_model = models[\"1WKDD+X6\"]\n",
    "    else:\n",
    "        raise\n",
    "    if args.evaluation_type.startswith(\"unify\"):\n",
    "        re_keys = [\"SameShape\", \"SameColor\", \"IsInside\"]\n",
    "        task_size = 6\n",
    "elif args.evaluation_type.startswith(\"parse\") or\\\n",
    "    args.evaluation_type.startswith(\"grounding\") or\\\n",
    "    args.evaluation_type.startswith(\"classify\") or\\\n",
    "    args.evaluation_type.startswith(\"generation\") or\\\n",
    "    args.evaluation_type.startswith(\"pc-parse+classify\") or\\\n",
    "    args.evaluation_type.startswith(\"yc-parse+classify\"):\n",
    "    assert args.color_avail == \"1,2\"\n",
    "    assert args.dataset.startswith(\"c-\") or args.dataset.startswith(\"pc-\") or args.dataset.startswith(\"yc-\") or args.dataset.startswith(\"u-\")\n",
    "    if args.evaluation_type.startswith(\"parse\"):\n",
    "        assert args.dataset == \"c-Lshape+Tshape+Cshape+Eshape+Fshape+Ashape\"\n",
    "    elif args.evaluation_type.startswith(\"pc-parse+classify\"):\n",
    "        assert args.dataset in [\"pc-Lshape+Tshape+Cshape+Eshape+Fshape+Ashape+Rect\",\n",
    "                                \"pc-Cshape+Eshape+Fshape+Ashape+Hshape+Rect\",\n",
    "                                \"pc-RectE1a+RectE2a+RectE3a\",\n",
    "                               ]\n",
    "        if args.dataset == \"pc-RectE1a+RectE2a+RectE3a\":\n",
    "            keys_dict = {\"Eshape\": 1, \"Rect\": 2}\n",
    "        else:\n",
    "            keys_dict = {\"Line\": 4}\n",
    "    elif args.evaluation_type.startswith(\"yc-parse+classify\"):\n",
    "        assert args.dataset in [\"yc-Eshape[5,9]+Fshape[5,9]+Ashape[5,9]\"]\n",
    "        if args.dataset == \"yc-Eshape[5,9]+Fshape[5,9]+Ashape[5,9]\":\n",
    "            keys_dict = {\"Line\": 4}\n",
    "        else:\n",
    "            raise\n",
    "    elif args.evaluation_type.startswith(\"classify\") or args.evaluation_type.startswith(\"generation\"):\n",
    "        assert args.dataset in [\"c-Eshape+Fshape+Ashape\", \"c-RectE1a+RectE2a+RectE3a\", \"u-graph-Graph1+Graph2+Graph3\"]\n",
    "        if args.dataset == \"c-RectE1a+RectE2a+RectE3a\":\n",
    "            assert args.canvas_size == 20\n",
    "        if args.dataset.startswith(\"u-\"):\n",
    "            assert args.canvas_size == 32\n",
    "    elif args.evaluation_type.startswith(\"grounding-RectE\"):\n",
    "        assert args.canvas_size == 20\n",
    "    if args.is_new_vertical:\n",
    "        re_keys = [\"Parallel\", \"VerticalMid\", \"VerticalEdge\"]\n",
    "    else:\n",
    "        re_keys = [\"Parallel\", \"Vertical\"]\n",
    "    task_size = 1\n",
    "else:\n",
    "    raise\n",
    "if args.concept_model_hash != \"None\":\n",
    "    models[args.concept_model_hash], info = load_model_hash(\n",
    "        args.concept_model_hash, return_args=True, isplot=0,\n",
    "        is_update_repr=True,\n",
    "        update_keys=[\n",
    "            \"Line\",\n",
    "            \"Lshape\",\n",
    "            \"Rect\",\n",
    "            \"RectSolid\",\n",
    "            \"Eshape\",\n",
    "            \"Fshape\",\n",
    "            \"Red\",\n",
    "            \"Blue\",\n",
    "            \"Green\",\n",
    "            \"Cube\",\n",
    "            \"Cylinder\",\n",
    "            \"Large\",\n",
    "            \"Small\",\n",
    "        ],\n",
    "        load_epoch=args.concept_load_id,\n",
    "    )\n",
    "    concept_model = models[args.concept_model_hash]\n",
    "    SGLD_args = info[\"args\"]\n",
    "    args.concept_load_id = info[\"load_id\"]\n",
    "else:\n",
    "    args.concept_load_id = \"None\"\n",
    "if args.relation_model_hash != \"None\":\n",
    "    models[args.relation_model_hash], info = load_model_hash(\n",
    "        args.relation_model_hash, return_args=True, isplot=0,\n",
    "        is_update_repr=True,\n",
    "        update_keys=[\n",
    "            \"SameShape\",\n",
    "            \"SameColor\",\n",
    "            \"IsInside\",\n",
    "            \"VerticalMid\",\n",
    "            \"VerticalEdge\",\n",
    "            \"Parallel\",\n",
    "            \"Vertical\",\n",
    "            \"IsEnclosed\",\n",
    "            \"IsNonOverlapXY\",\n",
    "            \"SameSize\",\n",
    "        ],\n",
    "        load_epoch=args.relation_load_id,\n",
    "    )\n",
    "    relation_model = models[args.relation_model_hash]\n",
    "    args.relation_load_id = info[\"load_id\"]\n",
    "else:\n",
    "    args.relation_load_id = \"None\"\n",
    "ebm_dict = Shared_Param_Dict(\n",
    "    concept_model=concept_model,\n",
    "    relation_model=relation_model,\n",
    ").to(device)\n",
    "for c_str in [\"Line\", \"Lshape\", \"Rect\", \"RectSolid\", \"Eshape\", \"Fshape\", \"Red\", \"Blue\", \"Green\", \"Cube\", \"Cylinder\", \"Large\", \"Small\"]:\n",
    "    ebm_dict.add_c_repr(CONCEPTS[c_str].get_node_repr()[None].to(device), c_str, ebm_mode=\"concept\")\n",
    "for c_str in [\"SameShape\", \"SameColor\", \"IsInside\", \"VerticalMid\", \"VerticalEdge\", \"Parallel\", \"Vertical\", \"SameSize\"]:\n",
    "    ebm_dict.add_c_repr(OPERATORS[c_str].get_node_repr()[None].to(device), c_str, ebm_mode=\"operator\")\n",
    "if args.evaluation_type.startswith(\"yc-\"):\n",
    "    if \"^\" in args.evaluation_type:\n",
    "        tasks = args.evaluation_type.split(\"^\")[-1].split(\"+\")\n",
    "    else:\n",
    "        tasks = [\"parse\", \"classify\"]\n",
    "\n",
    "    if \"classify\" in tasks:\n",
    "        models_3D = {}\n",
    "        print(\"Loading concept_3D:\")\n",
    "        models_3D[args.concept_model_hash_3D], info = load_model_hash(\n",
    "            args.concept_model_hash_3D, return_args=True, isplot=0,\n",
    "            is_update_repr=True, update_keys=[\"Line\", \"Lshape\", \"Rect\", \"RectSolid\", \"Eshape\", \"Fshape\"],\n",
    "            load_epoch=args.concept_load_id_3D,\n",
    "        )\n",
    "        concept_model_3D = models_3D[args.concept_model_hash_3D]\n",
    "        args.concept_load_id_3D = info[\"load_id\"]\n",
    "        print(\"Loading relation_3D:\")\n",
    "        models_3D[args.relation_model_hash_3D], info = load_model_hash(\n",
    "            args.relation_model_hash_3D, return_args=True, isplot=0,\n",
    "            is_update_repr=True, update_keys=[\"SameShape\", \"SameColor\", \"IsInside\", \"VerticalMid\", \"VerticalEdge\", \"Parallel\", \"Vertical\", \"IsEnclosed\", \"IsNonOverlapXY\"],\n",
    "            load_epoch=args.relation_load_id_3D,\n",
    "        )\n",
    "        relation_model_3D = models_3D[args.relation_model_hash_3D]\n",
    "        args.relation_load_id_3D = info[\"load_id\"]\n",
    "        ebm_dict_3D = Shared_Param_Dict(\n",
    "            concept_model=concept_model_3D,\n",
    "            relation_model=relation_model_3D,\n",
    "        ).to(device)\n",
    "        for c_str in [\"Line\", \"Lshape\", \"Rect\", \"RectSolid\", \"Eshape\", \"Fshape\"]:\n",
    "            ebm_dict_3D.add_c_repr(CONCEPTS[c_str].get_node_repr()[None].to(device), c_str, ebm_mode=\"concept\")\n",
    "        for c_str in [\"SameShape\", \"SameColor\", \"IsInside\", \"VerticalMid\", \"VerticalEdge\", \"Parallel\", \"Vertical\"]:\n",
    "            ebm_dict_3D.add_c_repr(OPERATORS[c_str].get_node_repr()[None].to(device), c_str, ebm_mode=\"operator\")\n",
    "\n",
    "# Composite dataset, canvas_size=8, up to 3b:\n",
    "if args.evaluation_type.startswith(\"yc-parse+classify\"):\n",
    "    composite_args = init_args({\n",
    "        \"dataset\": \"yc-Eshape[5,9]+Fshape[5,9]+Ashape[5,9]\",\n",
    "        \"seed_3d\": 42,\n",
    "        \"n_examples\": 200,\n",
    "        \"num_processes_3d\": 5,\n",
    "        \"n_queries_per_class\": 1,\n",
    "        \"color_map_3d\": \"same\",\n",
    "        \"add_thick_surf\": (0, 0.5),\n",
    "        \"add_thick_depth\": (0, 0.5),\n",
    "        \"image_size_3d\": (256, 256),\n",
    "        # 2D examples\n",
    "        \"seed\": 102,\n",
    "        \"use_seed_2d\": True,\n",
    "        \"canvas_size\": 16,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": \"1,2\",\n",
    "        \"max_n_distractors\": 0,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"parsing_check\": False,\n",
    "        \"allow_connect\": True,\n",
    "    })\n",
    "else:\n",
    "    composite_args = init_args({\n",
    "        \"dataset\": args.dataset,\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": args.val_n_examples,\n",
    "        \"canvas_size\": args.canvas_size,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": args.color_avail,\n",
    "        \"min_n_distractors\": 0,\n",
    "        \"max_n_distractors\": 2 if args.evaluation_type.startswith(\"grounding\") else 0,\n",
    "        \"allow_connect\": True,\n",
    "        \"parsing_check\": True if args.evaluation_type.startswith(\"grounding\") else False,\n",
    "    })\n",
    "    \n",
    "dataset, composite_args = get_dataset(composite_args, is_load=True)\n",
    "for key, value in get_SGLD_kwargs(SGLD_args).items():\n",
    "    if key not in [\"SGLD_mutual_exclusive_coef\",\n",
    "                   \"SGLD_pixel_entropy_coef\",\n",
    "                   \"SGLD_is_anneal\",\n",
    "                   \"SGLD_is_penalize_lower\",\n",
    "                   \"sample_step\",\n",
    "                   \"lambd\",\n",
    "                  ]:\n",
    "        setattr(args, key, value)\n",
    "args.is_two_branch = SGLD_args.is_two_branch\n",
    "args.is_image_tuple = False\n",
    "args.step_size_start = SGLD_args.step_size_start\n",
    "args.image_size = (args.canvas_size, args.canvas_size)\n",
    "args.SGLD_anneal_power = 2\n",
    "args.SGLD_fine_mutual_exclusive_coef = 0\n",
    "args.SGLD_mask_entropy_coef = 0\n",
    "args.rescaled_size = \"None\"\n",
    "if args.SGLD_is_penalize_lower_seq == \"None\":\n",
    "    args.SGLD_is_penalize_lower_seq = args.SGLD_is_penalize_lower\n",
    "filename = \"evaluation_{}_canvas_{}_color_{}_ex_{}_min_{}_model_{}_mutu_{}_ens_{}_sas_{}_newv_{}_batch_{}_con_{}_re_{}_bi_{}_range_{}_seed_{}_id_{}_Hash_{}_{}.p\".format(\n",
    "    args.evaluation_type, args.canvas_size, args.color_avail, args.val_n_examples, args.min_n_distractors, args.model_type, args.SGLD_mutual_exclusive_coef, args.ensemble_size, args.sample_step, args.is_new_vertical, args.val_batch_size, args.concept_model_hash, args.relation_model_hash, args.is_bidirectional_re, args.data_range, args.seed, args.id, get_hashing(str(args.__dict__), length=8), get_machine_name())\n",
    "print_banner(filename)\n",
    "pp.pprint(args.__dict__)\n",
    "data_record = {\"args\": args.__dict__}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38022597-9e32-4146-b579-45ba53e4be4f",
   "metadata": {},
   "source": [
    "# 3. Zero-shot tasks:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcc53e92-db22-43ad-b1d4-1a01ba30f280",
   "metadata": {},
   "source": [
    "### 3.1 Parse:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2522158d-09fb-47ae-92f5-108df5a6e979",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if args.evaluation_type.startswith(\"parse\"):\n",
    "    data_record[\"parsing_all_dict\"] = {}\n",
    "    set_seed(args.seed)\n",
    "    if args.model_type == \"rand-graph\":\n",
    "        freq_dict = {}\n",
    "        for data in dataset:\n",
    "            c_type_gt = get_c_type_from_data(data)\n",
    "            graph_gt = get_concept_graph(c_type_gt, is_new_vertical=args.is_new_vertical)\n",
    "            n_line = np.sum([1 for ele in graph_gt if ele[-1] == \"Line\"])\n",
    "            if args.is_new_vertical:\n",
    "                n_vertical_mid = np.sum([1 for ele in graph_gt if ele[-1] == 'VerticalMid'])\n",
    "                n_vertical_edge = np.sum([1 for ele in graph_gt if ele[-1] == 'VerticalEdge'])\n",
    "            else:\n",
    "                n_vertical = np.sum([1 for ele in graph_gt if ele[-1] == 'Vertical'])\n",
    "            n_parallel = np.sum([1 for ele in graph_gt if ele[-1] == 'Parallel'])\n",
    "            if args.is_new_vertical:\n",
    "                record_data(freq_dict, [n_line, n_vertical_mid, n_vertical_edge, n_parallel], [\"Line\", \"VerticalMid\", \"VerticalEdge\", \"Parallel\"])\n",
    "            else:\n",
    "                record_data(freq_dict, [n_line, n_vertical, n_parallel], [\"Line\", \"Vertical\", \"Parallel\"])\n",
    "        freq_mean = transform_dict(freq_dict, \"mean\")\n",
    "        if args.is_new_vertical:\n",
    "            freq_vertical_edge = freq_mean[\"VerticalEdge\"]\n",
    "            freq_vertical_mid = freq_mean[\"VerticalMid\"]\n",
    "        else:\n",
    "            freq_vertical = freq_mean[\"Vertical\"]\n",
    "        freq_parallel = freq_mean['Parallel']\n",
    "        if args.is_new_vertical:\n",
    "            freq_ver_para = np.array([freq_vertical_mid, freq_vertical_edge, freq_parallel])\n",
    "        else:\n",
    "            freq_ver_para = np.array([freq_vertical, freq_parallel])\n",
    "        freq_ver_para = freq_ver_para / freq_ver_para.sum()\n",
    "    else:\n",
    "        assert args.model_type == \"hc-ebm\"\n",
    "\n",
    "    make_dir(dirname + filename)\n",
    "    for i in range(len(dataset)):\n",
    "        data = dataset[i]\n",
    "        if i % args.inspect_interval == 0 and i > 0:\n",
    "            print(\"{}:\".format(i))\n",
    "            isplot = is_jupyter\n",
    "        else:\n",
    "            isplot = 0\n",
    "\n",
    "        if args.model_type == \"hc-ebm\":\n",
    "            input = data[0][None].to(device)\n",
    "            keys_dict = {\"Line\": 4}\n",
    "            selector = get_selector_for_parsing(keys_dict, ebm_dict, CONCEPTS, OPERATORS)\n",
    "            _, masks_top, info = get_selector_SGLD(\n",
    "                selector, input, args,\n",
    "                ebm_target=\"mask\",\n",
    "                init=args.init,\n",
    "                SGLD_object_exceed_coef=0,\n",
    "                SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "                SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "                ensemble_size=args.ensemble_size,\n",
    "                sample_step=args.sample_step,\n",
    "                isplot=isplot,\n",
    "            )\n",
    "            graph = []\n",
    "            for k, node_name in enumerate(selector.topological_sort):\n",
    "                c_type = node_name.split(\":\")[-1]\n",
    "                E_c = ebm_dict[c_type](input, (masks_top[k],)).item()\n",
    "                graph.append((k, c_type, E_c))\n",
    "\n",
    "            id_pairs = list(zip(*get_triu_ids(len(masks_top))))\n",
    "            for id0, id1 in id_pairs:\n",
    "                E_re = {}\n",
    "                for re_key in re_keys:\n",
    "                    if args.is_round_mask:\n",
    "                        E_re[re_key] = ebm_dict[re_key]((input, input), (masks_top[id0].round(), masks_top[id1].round())).item()\n",
    "                    else:\n",
    "                        E_re[re_key] = ebm_dict[re_key]((input, input), (masks_top[id0], masks_top[id1])).item()\n",
    "\n",
    "                idx_argmin = np.argmin(list(E_re.values()))\n",
    "                re_argmin = list(E_re)[idx_argmin]\n",
    "                E_min = list(E_re.values())[idx_argmin]\n",
    "                graph.append(((id0, id1), re_argmin, E_min))\n",
    "        elif args.model_type == \"rand-graph\":\n",
    "            graph = get_rand_graph(freq_dict, is_new_vertical=args.is_new_vertical)\n",
    "            input = None\n",
    "            masks_top = None\n",
    "        else:\n",
    "            raise\n",
    "        if i % args.inspect_interval == 0:\n",
    "            p.print(\"graph:\")\n",
    "            shortened_graph = [ele[:3] for ele in graph]\n",
    "            p.print(shortened_graph)\n",
    "            print(\"\\n\\n\")\n",
    "        record_data(data_record[\"parsing_all_dict\"],\n",
    "                    [i, graph, to_cpu_recur(masks_top), to_cpu_recur(input), data[3] if args.evaluation_type.startswith(\"parse\") else data[j][3]],\n",
    "                    [\"batch\", \"graph\", \"masks_top\", \"input\", \"info\"])\n",
    "        if i % args.inspect_interval == 0 and i > 0 or i == len(dataset) - 1:\n",
    "            pdump(data_record, dirname + filename)\n",
    "            p.print(\"Saved at {}\".format(dirname + filename))\n",
    "\n",
    "    if args.model_type == \"rand-graph\":\n",
    "        # Test if the rand-graph obeys the statistics:\n",
    "        graphs = data_record[\"parsing_all_dict\"][\"graph\"]\n",
    "        freq_dict2 = {}\n",
    "        for graph in graphs:\n",
    "            n_line = np.sum([1 for ele in graph if ele[-1] == \"Line\"])\n",
    "            n_vertical_edge = np.sum([1 for ele in graph if ele[-1] == 'VerticalEdge'])\n",
    "            n_vertical_mid = np.sum([1 for ele in graph if ele[-1] == 'VerticalMid'])\n",
    "            n_parallel = np.sum([1 for ele in graph if ele[-1] == 'Parallel'])\n",
    "            record_data(freq_dict2, [n_line, n_vertical_edge, n_vertical_mid, n_parallel], [\"Line\", \"VerticalEdge\", \"VerticalMid\", \"Parallel\"])\n",
    "        freq2_mean = transform_dict(freq_dict2, \"mean\")\n",
    "        freq2_vertical_edge = freq2_mean[\"VerticalEdge\"]\n",
    "        freq2_vertical_mid = freq2_mean[\"VerticalMid\"]\n",
    "        freq2_parallel = freq2_mean['Parallel']\n",
    "        freq2_ver_para = np.array([freq2_vertical_mid, freq2_vertical_edge, freq2_parallel])\n",
    "        freq2_ver_para = freq2_ver_para / freq2_ver_para.sum()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34d1efd5-bc03-4a95-a471-114925ad5aad",
   "metadata": {},
   "source": [
    "### 3.2 Grounding："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36fbbfa3-ab2b-4e71-a47c-df449794c01d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if args.evaluation_type.startswith(\"grounding\"):\n",
    "    isplot = is_jupyter\n",
    "    is_input = False\n",
    "    c_type = args.evaluation_type.split(\"-\")[-1]\n",
    "\n",
    "    graph = get_concept_graph(c_type,\n",
    "                              is_new_vertical=args.is_new_vertical,\n",
    "                              is_bidirectional_re=args.is_bidirectional_re,\n",
    "                              is_concept=args.is_concept,\n",
    "                              is_relation=args.is_relation,\n",
    "                             )\n",
    "    query = {\"graph\": graph}\n",
    "    set_seed(args.seed)\n",
    "    allow_connect = False\n",
    "    if c_type == \"Eshape\":\n",
    "        if args.is_proper_size:\n",
    "            dataset_type = \"c-Eshape[6,8]+Cshape+Lshape+Tshape+Rect+RectSolid^Eshape[6,8]\"\n",
    "        else:\n",
    "            dataset_type = \"c-Eshape+Cshape+Lshape+Tshape+Rect+RectSolid^Eshape\"\n",
    "    elif c_type == \"Fshape\":\n",
    "        if args.is_proper_size:\n",
    "            dataset_type = \"c-Fshape[6,8]+Cshape+Lshape+Tshape+Rect+RectSolid^Fshape[6,8]\"\n",
    "        else:\n",
    "            dataset_type = \"c-Fshape+Cshape+Lshape+Tshape+Rect+RectSolid^Fshape\"\n",
    "    elif c_type == \"Ashape\":\n",
    "        if args.is_proper_size:\n",
    "            dataset_type = \"c-Ashape[6,8]+Cshape+Lshape+Tshape+Rect+RectSolid^Ashape[6,8]\"\n",
    "        else:\n",
    "            dataset_type = \"c-Ashape+Cshape+Lshape+Tshape+Rect+RectSolid^Ashape\"\n",
    "    elif c_type == \"RectE1a\":\n",
    "        if args.is_harder_distractor:\n",
    "            dataset_type = \"c-RectE1a+Eshape+Rect+Tshape+Fshape+Ashape^RectE1a\"\n",
    "            allow_connect = True\n",
    "        else:\n",
    "            dataset_type = \"c-RectE1a+Cshape+Lshape+Tshape+Rect+RectSolid^RectE1a\"\n",
    "    elif c_type == \"RectE2a\":\n",
    "        if args.is_harder_distractor:\n",
    "            dataset_type = \"c-RectE2a+Eshape+Rect+Tshape+Fshape+Ashape^RectE2a\"\n",
    "            allow_connect = True\n",
    "        else:\n",
    "            dataset_type = \"c-RectE2a+Cshape+Lshape+Tshape+Rect+RectSolid^RectE2a\"\n",
    "    elif c_type == \"RectE3a\":\n",
    "        if args.is_harder_distractor:\n",
    "            dataset_type = \"c-RectE3a+Eshape+Rect+Tshape+Fshape+Ashape^RectE3a\"\n",
    "            allow_connect = True\n",
    "        else:\n",
    "            dataset_type = \"c-RectE3a+Cshape+Lshape+Tshape+Rect+RectSolid^RectE3a\"\n",
    "    else:\n",
    "        raise\n",
    "    composite_args = init_args({\n",
    "        \"dataset\": dataset_type,\n",
    "        \"seed\": 2,\n",
    "        \"n_examples\": args.val_n_examples,\n",
    "        \"canvas_size\": args.canvas_size,\n",
    "        \"rainbow_prob\": 0.,\n",
    "        \"w_type\": \"image+mask\",\n",
    "        \"color_avail\": args.color_avail,\n",
    "        \"min_n_distractors\": args.min_n_distractors,\n",
    "        \"max_n_distractors\": args.max_n_distractors,\n",
    "        \"allow_connect\": allow_connect,\n",
    "        \"parsing_check\": True,\n",
    "    })\n",
    "    dataset, composite_args = get_dataset(composite_args, is_load=True)\n",
    "    selector = get_selector_from_graph(query[\"graph\"], ebm_dict, CONCEPTS, OPERATORS)\n",
    "    Dict = {\"args\": args.__dict__}\n",
    "    dataloader = DataLoader(dataset[:200], batch_size=args.val_batch_size, collate_fn=Batch(is_collate_tuple=True).collate(), shuffle=False)\n",
    "    make_dir(dirname + filename)\n",
    "    for i, data in enumerate(dataloader):\n",
    "        if isplot:\n",
    "            dataset.draw(i)\n",
    "        c_type = get_c_type_from_data(data)\n",
    "        img = data[0].to(device)\n",
    "        graph = get_concept_graph(c_type,\n",
    "                                  is_new_vertical=args.is_new_vertical,\n",
    "                                  is_bidirectional_re=args.is_bidirectional_re,\n",
    "                                  is_concept=args.is_concept,\n",
    "                                  is_relation=args.is_relation,\n",
    "                                 )\n",
    "        query = {\"graph\": graph}\n",
    "        if args.model_type == \"hc-ebm\":\n",
    "            selector = get_selector_from_graph(query[\"graph\"], ebm_dict, CONCEPTS, OPERATORS)\n",
    "            _, masks_top, info = get_selector_SGLD(\n",
    "                selector, img, args,\n",
    "                ebm_target=\"mask\",\n",
    "                init=args.init,\n",
    "                SGLD_object_exceed_coef=0,\n",
    "                SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "                SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "                ensemble_size=args.ensemble_size,\n",
    "                sample_step=args.sample_step,\n",
    "                isplot=isplot,\n",
    "            )  # Each element of masks_top: [topk:1, B, 1, H, W]\n",
    "            mask_composite = torch.stack(masks_top).max(0)[0][0]\n",
    "            if isplot:\n",
    "                for i in range(len(info[\"mask_list\"])):\n",
    "                    print(\"{} th mask:\".format(i))\n",
    "                    plot_matrices(info[\"mask_list\"][i].squeeze()[::2], images_per_row=15, scale_limit=(0,1))\n",
    "                print(\"max mask:\")\n",
    "                plot_matrices(np.stack(info[\"mask_list\"]).max(0).squeeze()[::2], images_per_row=15, scale_limit=(0,1))\n",
    "            if is_input:\n",
    "                user_input = input(\"continue [y/n]:\")\n",
    "                if user_input in [\"y\", \"Y\"]:\n",
    "                    pass\n",
    "                else:\n",
    "                    raise\n",
    "        elif args.model_type == \"rand-obj\":\n",
    "            mask_composite = random_obj(img)\n",
    "        iou = to_np_array(get_soft_IoU(mask_composite.round(), data[1][0].round().to(device), dim=(-3,-2,-1)), full_reduce=False)\n",
    "        record_data(Dict, [iou], [\"iou\"])\n",
    "        record_data(Dict, [img, np.stack(to_np_array(*masks_top, keep_list=True), 2)[0]], [\"input\", \"masks_top\"])\n",
    "        if isplot:\n",
    "            plot_matrices([mask_composite.squeeze(0)[0], data[1][0].round()[0,0]], images_per_row=6,\n",
    "                          subtitles=[\"iou: {:.3f}\".format(to_np_array(iou)), None])\n",
    "        if i % args.inspect_interval == 0 or i == len(dataset) - 1:\n",
    "            p.print(\"{}: iou_mean {:.4f}\".format(i, np.mean(np.concatenate(Dict[\"iou\"]))))\n",
    "            try_call(pdump, args=[Dict, dirname + filename], max_exp_time=300)\n",
    "    Dict[\"iou\"] = np.concatenate(Dict[\"iou\"])\n",
    "    pdump(Dict, dirname + filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02c08263-0a96-4094-9fec-a23196b1db35",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# input = data[0].to(device)[None]\n",
    "# graph = get_concept_graph(c_type,\n",
    "#                           is_new_vertical=args.is_new_vertical,\n",
    "#                           is_bidirectional_re=args.is_bidirectional_re,\n",
    "#                           is_concept=args.is_concept,\n",
    "#                          )\n",
    "# query = {\"graph\": graph}\n",
    "# if args.model_type == \"hc-ebm\":\n",
    "#     selector = get_selector_from_graph(query[\"graph\"], ebm_dict, CONCEPTS, OPERATORS)\n",
    "#     _, masks_top, info = get_selector_SGLD(\n",
    "#         selector, input, args,\n",
    "#         ebm_target=\"mask\",\n",
    "#         init=args.init,\n",
    "#         SGLD_object_exceed_coef=0,\n",
    "#         SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "#         SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "#         ensemble_size=args.ensemble_size,\n",
    "#         sample_step=args.sample_step,\n",
    "#         isplot=isplot,\n",
    "#     )  # Each element of masks_top: [topk:1, B, 1, H, W]\n",
    "#     mask_composite = torch.stack(masks_top).max(0)[0][0]\n",
    "#     plot_matrices([mask_composite.squeeze(0)[0], data[1][0].round()[0]], images_per_row=6,\n",
    "#                    subtitles=[\"iou: {:.3f}\".format(to_np_array(iou)), None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f99c0229-845f-49cb-8d59-ea6659a6e0c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mask2 = torch.zeros(16,16).to(device)\n",
    "# mask2[5:13, -1:] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d7e1f39-3a18-4f93-b182-e332b104a3ef",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# mask4 = torch.zeros(16,16).to(device)\n",
    "# mask4[-4, -6:] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecf37a62-fe5c-4a18-ace6-035cd441ad78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mask1 = torch.zeros(16,16).to(device)\n",
    "# mask1[5, -6:] = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5417a415-d0f0-412c-b6f6-34f395c8840c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# masks_top2= (mask2[None,None],mask1[None,None], masks_top[2][0], mask4[None,None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0293ea34-3fe1-4204-a8fb-e09ca88ee323",
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot_matrices(tuple(mask.squeeze() for mask in masks_top2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d792497d-f24d-45d9-a8f5-543b73fba963",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if args.is_analysis:\n",
    "#     c_type = \"Parallel\"\n",
    "#     graph = get_concept_graph(c_type)\n",
    "#     selector = get_selector_from_graph(graph, ebm_dict, CONCEPTS, OPERATORS)\n",
    "\n",
    "#     col = 4\n",
    "#     col2 = 10\n",
    "#     image = torch.zeros(16,16).to(device)\n",
    "#     image[4:12, col:col+1] = 1\n",
    "#     image[10:12, col2:col2+1] = 1\n",
    "#     mask1 = torch.zeros(16, 16).to(device)\n",
    "#     mask1[4:12, col:col+1] = 1\n",
    "#     mask2 = torch.zeros(16, 16).to(device)\n",
    "#     mask2[10:12, col2:col2+1] = 1\n",
    "\n",
    "#     img = to_one_hot(image).to(device)\n",
    "#     E = selector(img[None], (mask1[None,None], mask2[None,None]))\n",
    "#     visualize_matrices([image], subtitles=[\"E={:.4f}\".format(to_np_array(E.item()))])\n",
    "#     plot_matrices([mask1, mask2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60b02673-ccd3-49bf-9452-138668c8186b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# c_type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7305222c-41fd-4201-810a-86bfec3075d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if args.is_analysis:\n",
    "#     c_type = \"VerticalEdge\"\n",
    "#     graph = get_concept_graph(c_type)\n",
    "#     selector = get_selector_from_graph(graph, ebm_dict, CONCEPTS, OPERATORS)\n",
    "#     row = 15\n",
    "#     col = 15\n",
    "#     image = torch.zeros(16,16).to(device)\n",
    "#     image[4:row, col:col+1] = 1\n",
    "#     image[row:row+1, 7:15] = 1\n",
    "#     mask1 = torch.zeros(16, 16).to(device)\n",
    "#     mask1[4:row, col:col+1] = 1\n",
    "#     mask2 = torch.zeros(16, 16).to(device)\n",
    "#     mask2[row:row+1, 7:15] = 1\n",
    "\n",
    "#     img = to_one_hot(image).to(device)\n",
    "#     E = selector(img[None], (mask1[None,None], mask2[None,None]))\n",
    "#     visualize_matrices([image], subtitles=[\"E={:.4f}\".format(to_np_array(E.item()))])\n",
    "#     plot_matrices([mask1, mask2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "425a1851-3d24-49aa-804b-2e20753f28b4",
   "metadata": {},
   "source": [
    "### 3.3 Classify:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da78aa7-aea5-4822-acab-7aa1da56fe5d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if args.evaluation_type.startswith(\"classify\"):\n",
    "    set_seed(args.seed)\n",
    "    if args.data_range == \"None\":\n",
    "        include_ids = None\n",
    "    else:\n",
    "        include_ids = list(np.arange(*([int(ele) for ele in args.data_range.split(\":\")])))\n",
    "    isplot = is_jupyter\n",
    "    if args.dataset.startswith(\"u-\"):\n",
    "        c_types = args.dataset.split(\"-\")[-1].split(\"+\")\n",
    "    else:\n",
    "        c_types = args.dataset.split(\"-\")[1].split(\"+\")\n",
    "    c_type_graphs = {c_type: get_concept_graph(c_type,\n",
    "                                               is_new_vertical=args.is_new_vertical,\n",
    "                                               is_bidirectional_re=args.is_bidirectional_re,\n",
    "                                               is_concept=args.is_concept,\n",
    "                                               is_relation=args.is_relation,\n",
    "                                              ) for c_type in c_types}\n",
    "    selectors = {c_type: get_selector_from_graph(graph, ebm_dict, CONCEPTS, OPERATORS) for c_type, graph in c_type_graphs.items()}\n",
    "    Dict = {\"args\": args.__dict__, \"results\": {}}\n",
    "    dataloader = DataLoader(dataset[:200], batch_size=args.val_batch_size, collate_fn=Batch(is_collate_tuple=True).collate(), shuffle=False)\n",
    "    alpha_list = np.linspace(0,1,21).round(3)\n",
    "    make_dir(dirname + filename)\n",
    "    for i, data in enumerate(dataloader):\n",
    "        if include_ids is not None:\n",
    "            if i not in include_ids:\n",
    "                print(f\"Skip: {i}\")\n",
    "                continue\n",
    "        if i % args.inspect_interval == 0:\n",
    "            p.print(\"iter: {}\".format(i))\n",
    "        energy_dict = {}\n",
    "        for key, selector in selectors.items():\n",
    "            img = data[0].to(device)\n",
    "            if args.dataset.startswith(\"u-\"):\n",
    "                img = F.interpolate(img, size=(32,32), mode=\"nearest\")\n",
    "            if isplot:\n",
    "                if data[2][0] == key:\n",
    "                    p.print(\"{}: [{}]:\".format(i, key), banner_size=50)\n",
    "                else:\n",
    "                    p.print(\"{}: {}:\".format(i, key), banner_size=50)\n",
    "            _, masks_top, info = get_selector_SGLD(\n",
    "                selector, img, args,\n",
    "                ebm_target=\"mask\",\n",
    "                init=args.init,\n",
    "                SGLD_object_exceed_coef=0,\n",
    "                SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "                SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "                ensemble_size=args.ensemble_size,\n",
    "                sample_step=args.sample_step,\n",
    "                isplot=isplot,\n",
    "            )  # mask_top: each element [topk, B, 1, H, W]\n",
    "            energy_dict[key] = {}\n",
    "            energy_dict[key][\"masks_top\"] = np.stack(to_np_array(*masks_top, keep_list=True), 2)[0]  # [B, n_masks, 1, H, W]\n",
    "            energy_dict[key][\"energy_mask\"] = np.stack(list(info['energy_mask'].values()), 2)[0]  # [B, n_masks_concept_re]\n",
    "            energy_dict[key][\"energy_mean\"] = np.mean(energy_dict[key][\"energy_mask\"], -1)  # [B]\n",
    "            energy_dict[key][\"energy_sum\"] = np.sum(energy_dict[key][\"energy_mask\"], -1)  # [B]\n",
    "            energy_dict[key][\"mutual_exclusive\"] = info['mutual_exclusive_list'][0]\n",
    "            for alpha in alpha_list:\n",
    "                energy_dict[key][\"energy_total^m:{:.2f}\".format(alpha)] = energy_dict[key][\"energy_mean\"] + energy_dict[key][\"mutual_exclusive\"] * alpha\n",
    "                energy_dict[key][\"energy_total^s:{:.2f}\".format(alpha)] = energy_dict[key][\"energy_sum\"] + energy_dict[key][\"mutual_exclusive\"] * alpha\n",
    "            \n",
    "            gt_mask = to_np_array(img[:,:1] != 1)\n",
    "            masks_top_sum = np.stack(to_np_array(*masks_top, keep_list=True)).sum(0)[0]\n",
    "            energy_dict[key][\"pixels_under\"] = np.clip(gt_mask - masks_top_sum, a_min=0, a_max=None).sum((-3,-2,-1))\n",
    "            energy_dict[key][\"pixels_exceed_out\"] = np.clip((masks_top_sum - gt_mask) * (1-gt_mask), a_min=0, a_max=None).sum((-3,-2,-1))\n",
    "            batch_size = masks_top_sum.shape[0]\n",
    "            energy_dict[key][\"pixels_exceed_out_mean\"] = np.clip((masks_top_sum - gt_mask).reshape(batch_size, -1)[:, (1-gt_mask).reshape(-1).astype(bool)], a_min=0, a_max=None).mean((-1))\n",
    "            energy_dict[key][\"pixels_exceed_in\"] = np.clip((masks_top_sum - gt_mask) * gt_mask, a_min=0, a_max=None).sum((-3,-2,-1))\n",
    "\n",
    "            for subkey in energy_dict[key]:\n",
    "                record_data(Dict[\"results\"], energy_dict[key][subkey], \"{}:{}\".format(key, subkey))\n",
    "            if isplot:\n",
    "                for ii in range(len(info[\"mask_list\"])):\n",
    "                    print(\"{} th mask:\".format(ii))\n",
    "                    plot_matrices(info[\"mask_list\"][ii].squeeze()[::2], images_per_row=15, scale_limit=(0,1))\n",
    "                print(\"max mask:\")\n",
    "                plot_matrices(np.stack(info[\"mask_list\"]).max(0).squeeze()[::2], images_per_row=15, scale_limit=(0,1))\n",
    "            if is_input:\n",
    "                user_input = input(\"continue [y/n]:\")\n",
    "                if user_input in [\"y\", \"Y\"]:\n",
    "                    pass\n",
    "                else:\n",
    "                    raise\n",
    "        record_data(Dict[\"results\"], np.array(data[2]), \"ground_truth\")\n",
    "        for alpha in alpha_list:\n",
    "            energy_total_list_m = []\n",
    "            energy_total_list_s = []\n",
    "            for c_type in c_types:\n",
    "                energy_total_list_m.append(energy_dict[c_type][\"energy_total^m:{:.2f}\".format(alpha)])\n",
    "                energy_total_list_s.append(energy_dict[c_type][\"energy_total^s:{:.2f}\".format(alpha)])\n",
    "            energy_total_list_m = np.stack(energy_total_list_m, -1)  # [B, n_c_types]\n",
    "            energy_total_list_s = np.stack(energy_total_list_s, -1)  # [B, n_c_types]\n",
    "            pred_argmin_m = np.argmin(energy_total_list_m, -1)  # [B]\n",
    "            pred_argmin_s = np.argmin(energy_total_list_s, -1)  # [B]\n",
    "            pred_m = np.array([c_types[k] for k in pred_argmin_m])\n",
    "            pred_s = np.array([c_types[k] for k in pred_argmin_s])\n",
    "            record_data(Dict[\"results\"], pred_m, \"pred^m:{:.2f}\".format(alpha))\n",
    "            record_data(Dict[\"results\"], pred_s, \"pred^s:{:.2f}\".format(alpha))\n",
    "            record_data(Dict[\"results\"], Dict[\"results\"][\"pred^m:{:.2f}\".format(alpha)][-1] == Dict[\"results\"][\"ground_truth\"][-1], \"acc^m:{:.2f}\".format(alpha))\n",
    "            record_data(Dict[\"results\"], Dict[\"results\"][\"pred^s:{:.2f}\".format(alpha)][-1] == Dict[\"results\"][\"ground_truth\"][-1], \"acc^s:{:.2f}\".format(alpha))\n",
    "        if i % args.inspect_interval == 0 or i == len(dataloader) - 1:\n",
    "            p.print(\"gt: {}\".format(data[2][0]))\n",
    "            for alpha in alpha_list:\n",
    "                print(\"{}: \\tpred^m: {} \\tacc^m: {:.4f}\\t acc_mean^m: {:.4f}\".format(\n",
    "                    i,\n",
    "                    Dict[\"results\"][\"pred^m:{:.2f}\".format(alpha)][-1][0],\n",
    "                    Dict[\"results\"][\"acc^m:{:.2f}\".format(alpha)][-1][0],\n",
    "                    np.concatenate(Dict[\"results\"][\"acc^m:{:.2f}\".format(alpha)]).mean(),\n",
    "                ))\n",
    "            print()\n",
    "            for alpha in alpha_list:\n",
    "                print(\"{}: \\tpred^s: {} \\tacc^s: {:.4f}\\t acc_mean^s: {:.4f}\".format(\n",
    "                    i,\n",
    "                    Dict[\"results\"][\"pred^s:{:.2f}\".format(alpha)][-1][0],\n",
    "                    Dict[\"results\"][\"acc^s:{:.2f}\".format(alpha)][-1][0],\n",
    "                    np.concatenate(Dict[\"results\"][\"acc^s:{:.2f}\".format(alpha)]).mean(),\n",
    "                ))\n",
    "            try_call(pdump, args=[Dict, dirname + filename], max_exp_time=300)\n",
    "    Dict[\"results\"] = transform_dict(Dict[\"results\"], \"concatenate\")\n",
    "    pdump(Dict, dirname + filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f062b30a-6abe-4ef3-8689-b770b08407d6",
   "metadata": {},
   "source": [
    "### 3.4 Parse2D+classify3D:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48869307-7332-4720-bc3b-5d0bb4731247",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if args.evaluation_type.startswith(\"yc\"):\n",
    "    set_seed(args.seed)\n",
    "    alpha_list = np.linspace(0, 1, 21).round(2)\n",
    "    isplot = is_jupyter\n",
    "    data_record = {\"args\": args.__dict__}\n",
    "    data_record[\"results\"] = {}\n",
    "    dataloader = DataLoader(dataset, batch_size=1, collate_fn=Batch(is_collate_tuple=True).collate(), shuffle=False)\n",
    "\n",
    "    if \"parse\" in tasks:\n",
    "        for i, data in enumerate(dataloader):\n",
    "            inputs = data[0][0]\n",
    "            concept_labels = [ele[0] for ele in data[0][2]]\n",
    "            example_labels = [ele[0] for ele in data[1][2]]\n",
    "            example_idx = np.array([concept_labels.index(example_label) for example_label in example_labels])\n",
    "            record_data(data_record, [concept_labels, example_labels, example_idx], [\"concept_labels\", \"example_labels\", \"example_idx\"])\n",
    "            if i % args.inspect_interval == 0 or i == len(dataloader) - 1:\n",
    "                p.print(\"Task {}:\".format(i), banner_size=100)\n",
    "                print(\"concept_labels: {}\".format(concept_labels))\n",
    "                if \"^parse\" not in args.evaluation_type:\n",
    "                    print(\"example_labels: {}\".format(example_labels))\n",
    "                    print(\"example_idx: {}\\n\".format(example_idx))\n",
    "\n",
    "            distance_edit_dict_all = {}\n",
    "            is_isomorphic_dict_all = {}\n",
    "            for j, input in enumerate(inputs):\n",
    "                input = input.to(device)\n",
    "                if (i % args.inspect_interval == 0 or i == len(dataloader) - 1) and args.verbose >= 1:\n",
    "                    print(\"Task {}, parsing concept {}, type {}:\".format(i, j, data[0][2][j][0]))\n",
    "                graph_trimmed_dict, graph_dict, masks_top, info = parse_selector_from_image(\n",
    "                    input=input,\n",
    "                    args=args,\n",
    "                    keys_dict=keys_dict,\n",
    "                    isplot=isplot,\n",
    "                    infer_order=args.infer_order,\n",
    "                    topk=args.topk,\n",
    "                    init=args.init,\n",
    "                )\n",
    "                graph_gt = get_concept_graph(concept_labels[j], is_new_vertical=args.is_new_vertical)\n",
    "                distance_edit_dict = {k: get_graph_edit_distance(graph_trimmed_dict[k], graph_gt, to_undirected=True) for k in range(args.topk)}\n",
    "                is_isomorphic_dict = {k: distance_edit_dict[k] == 0 for k in range(args.topk)}\n",
    "                distance_edit_dict_all[concept_labels[j]] = distance_edit_dict\n",
    "                is_isomorphic_dict_all[concept_labels[j]] = is_isomorphic_dict\n",
    "                masks_top_sum = np.stack(to_np_array(*masks_top, keep_list=True)).sum(0)\n",
    "                gt_mask = to_np_array(input[:,:1] != 1)\n",
    "                # if is_isomorphic_dict_all[concept_labels[j]][0] != 1:\n",
    "                #     print(\"error:\")\n",
    "                #     pp.pprint(graph_gt)\n",
    "                #     pp.pprint(graph_trimmed_dict[0])\n",
    "                #     isplot = True\n",
    "                #     print()\n",
    "                # else:\n",
    "                #     isplot = False\n",
    "\n",
    "                pixels_under = np.clip(gt_mask - masks_top_sum, a_min=0, a_max=None).sum((-3,-2,-1))\n",
    "                pixels_exceed_out = np.clip((masks_top_sum - gt_mask) * (1-gt_mask), a_min=0, a_max=None).sum((-3,-2,-1))\n",
    "                pixels_exceed_out_mean = np.clip((masks_top_sum - gt_mask).reshape(masks_top_sum.shape[-1], -1)[:, (1-gt_mask).flatten().astype(bool)], a_min=0, a_max=None).mean((-1))\n",
    "                pixels_exceed_in = np.clip((masks_top_sum - gt_mask) * gt_mask, a_min=0, a_max=None).sum((-3,-2,-1))\n",
    "                record_data(data_record,\n",
    "                            [graph_gt, np.stack(to_np_array(*masks_top, keep_list=True)), masks_top_sum, gt_mask, pixels_under, pixels_exceed_out, pixels_exceed_out_mean, pixels_exceed_in, info],\n",
    "                            [f\"graph_gt_{concept_labels[j]}\", f\"masks_top_{concept_labels[j]}\", f\"masks_top_sum_{concept_labels[j]}\", f\"gt_mask_{concept_labels[j]}\", f\"pixels_under_{concept_labels[j]}\", f\"pixels_exceed_out_{concept_labels[j]}\", f\"pixels_exceed_out_mean_{concept_labels[j]}\", f\"pixels_exceed_in_{concept_labels[j]}\", f\"info_{concept_labels[j]}\"])\n",
    "                record_data(data_record,\n",
    "                            [graph_gt, np.stack(to_np_array(*masks_top, keep_list=True)), masks_top_sum, gt_mask, pixels_under, pixels_exceed_out, pixels_exceed_out_mean, pixels_exceed_in, info],\n",
    "                            [f\"graph_gt_ex_{j}\", f\"masks_top_ex_{j}\", f\"masks_top_sum_ex_{j}\", f\"gt_mask_ex_{j}\", f\"pixels_under_ex_{j}\", f\"pixels_exceed_out_ex_{j}\", f\"pixels_exceed_out_mean_ex_{j}\", f\"pixels_exceed_in_ex_{j}\", f\"info_ex_{j}\"])\n",
    "        \n",
    "                energy_list = []\n",
    "                concept_energy_list = []\n",
    "                relation_energy_list = []\n",
    "                for k in range(args.topk):\n",
    "                    energy_list.append(np.sum([ele[2] for ele in graph_trimmed_dict[k]]))\n",
    "                    concept_energy_list.append(np.sum([ele[2] for ele in graph_trimmed_dict[k] if not isinstance(ele[0], tuple)]))\n",
    "                    relation_energy_list.append(np.sum([ele[2] for ele in graph_trimmed_dict[k] if isinstance(ele[0], tuple)]))\n",
    "                    record_data(data_record, [graph_trimmed_dict[k], graph_dict[k], distance_edit_dict[k], is_isomorphic_dict[k], energy_list[-1], concept_energy_list[-1], relation_energy_list[-1]],\n",
    "                                [f\"graph_trimmed_{concept_labels[j]}_{k}\", f\"graph_{concept_labels[j]}_{k}\", f\"distance_edit_{concept_labels[j]}_{k}\", f\"is_isomorphic_{concept_labels[j]}_{k}\", f\"energy_{concept_labels[j]}_{k}\", f\"concept_energy_{concept_labels[j]}_{k}\", f\"relation_energy_{concept_labels[j]}_{k}\"])\n",
    "                    record_data(data_record, [graph_trimmed_dict[k], graph_dict[k], distance_edit_dict[k], is_isomorphic_dict[k], energy_list[-1], concept_energy_list[-1], relation_energy_list[-1]],\n",
    "                                [f\"graph_trimmed_ex_{j}_{k}\", f\"graph_ex_{j}_{k}\", f\"distance_edit_ex_{j}_{k}\", f\"is_isomorphic_ex_{j}_{k}\", f\"energy_ex_{j}_{k}\", f\"concept_energy_ex_{j}_{k}\", f\"relation_energy_ex_{j}_{k}\"])\n",
    "                energy_list = np.array(energy_list)\n",
    "                concept_energy_list = np.array(concept_energy_list)\n",
    "                relation_energy_list = np.array(relation_energy_list)\n",
    "                record_data(data_record, [energy_list, concept_energy_list, relation_energy_list],\n",
    "                            [f\"energy_{concept_labels[j]}\", f\"concept_energy_{concept_labels[j]}\", f\"relation_energy_{concept_labels[j]}\"])\n",
    "                record_data(data_record, [energy_list, concept_energy_list, relation_energy_list],\n",
    "                            [f\"energy_ex_{j}\", f\"concept_energy_ex_{j}\", f\"relation_energy_ex_{j}\"])\n",
    "                record_data(data_record, [np.mean(list(distance_edit_dict.values())), np.mean(list(is_isomorphic_dict.values()))],\n",
    "                            [f\"distance_edit_{concept_labels[j]}\", f\"is_isomorphic_{concept_labels[j]}\"])\n",
    "                record_data(data_record, [np.mean(list(distance_edit_dict.values())), np.mean(list(is_isomorphic_dict.values()))],\n",
    "                            [f\"distance_edit_ex_{j}\", f\"is_isomorphic_ex_{j}\"])\n",
    "                pairwise_distance_topk = np.array([[get_graph_edit_distance(graph_trimmed_dict[ll], graph_trimmed_dict[mm], to_undirected=True) for mm in range(args.topk)] for ll in range(args.topk)])\n",
    "                record_data(data_record, [pairwise_distance_topk], [f\"pairwise_distance_topk_{concept_labels[j]}\"])\n",
    "                record_data(data_record, [pairwise_distance_topk], [f\"pairwise_distance_topk_ex_{j}\"])\n",
    "                if (i % args.inspect_interval == 0 or i == len(dataloader) - 1) and args.verbose >= 1:\n",
    "                    print(\"distance_edit:\")\n",
    "                    print(np.array(list(distance_edit_dict.values())))\n",
    "                    print(\"energy:\")\n",
    "                    print(energy_list)\n",
    "                    print(\"concept_energy:\")\n",
    "                    print(concept_energy_list)\n",
    "                    print(\"relation_energy:\")\n",
    "                    print(relation_energy_list)\n",
    "                    print(\"pixels_under:\")\n",
    "                    print(pixels_under)\n",
    "                    print(\"pixels_exceed_in:\")\n",
    "                    print(pixels_exceed_in)\n",
    "                    print(\"pixels_exceed_out:\")\n",
    "                    print(pixels_exceed_out)\n",
    "                    print(\"pixels_exceed_out_mean:\")\n",
    "                    print(pixels_exceed_out_mean)\n",
    "                    print(f\"pairwise_distance_topk_{concept_labels[j]}:\")\n",
    "                    print(pairwise_distance_topk.mean(-1))\n",
    "                    print()\n",
    "                del input\n",
    "                del masks_top\n",
    "                gc.collect()\n",
    "                if (i % args.inspect_interval == 0 or i == len(dataloader) - 1) and args.verbose >= 2:\n",
    "                    print(\"{}th concept, gt concept: {}\".format(j, concept_labels[j]))\n",
    "                    print(\"graph_gt:\")\n",
    "                    pp.pprint(graph_gt)\n",
    "                    print(\"graph_trimmed:\")\n",
    "                    pp.pprint(graph_trimmed_dict)\n",
    "                    print(\"Edit distance: {}   is_isomorphic: {}\".format(distance_edit_dict, is_isomorphic_dict))\n",
    "                    print()\n",
    "                    if isplot:\n",
    "                        for i in range(1):\n",
    "                            print(f\"top {i} SGLD:\")\n",
    "                            for ll in range(len(info[\"mask_list\"])):\n",
    "                                print(f\"  mask {ll}:\")\n",
    "                                plot_matrices(info[\"mask_list\"][ll][::2,i].squeeze((1,2)), images_per_row=19, figsize=(20,4))\n",
    "                            print(\"Max mask:\")\n",
    "                            plot_matrices(np.stack(info[\"mask_list\"]).max(0)[::2,i].squeeze((1,2)), images_per_row=19, figsize=(20,4))\n",
    "            for k in range(args.topk):\n",
    "                record_data(data_record, [np.mean([distance_edit_dict_all[c_type][k] for c_type in concept_labels]), np.mean([is_isomorphic_dict_all[c_type][k] for c_type in concept_labels])],\n",
    "                            [f\"distance_edit_{k}\", f\"is_isomorphic_{k}\"])\n",
    "            record_data(data_record, [np.mean([distance_edit_dict_all[c_type][k] for k in range(args.topk) for c_type in concept_labels]), np.mean([is_isomorphic_dict_all[c_type][k] for k in range(args.topk) for c_type in concept_labels])],\n",
    "                            [f\"distance_edit\", f\"is_isomorphic\"])\n",
    "            for c_type in sorted(concept_labels):\n",
    "                for k in range(args.topk):\n",
    "                    print(\"{}_{}:     \\tedit_distance_mean: {:.4f}    acc_iso_mean: {:.4f}\".format(\n",
    "                        c_type, k, np.mean(data_record[f\"distance_edit_{c_type}_{k}\"]), np.mean(data_record[f\"is_isomorphic_{c_type}_{k}\"])))\n",
    "                print(\"{}_m:     \\tedit_distance_mean: {:.4f}    acc_iso_mean: {:.4f}\".format(\n",
    "                    c_type, np.mean([np.mean(data_record[f\"distance_edit_{c_type}_{k}\"]) for k in range(args.topk)]),\n",
    "                    np.mean([np.mean(data_record[f\"is_isomorphic_{c_type}_{k}\"]) for k in range(args.topk)])))\n",
    "            print()\n",
    "            for c_type in sorted(concept_labels):\n",
    "                print(\"{}_m:     \\tedit_distance_mean: {:.4f}    acc_iso_mean: {:.4f}\".format(\n",
    "                    c_type, np.mean(data_record[f\"distance_edit_{c_type}\"]), np.mean(data_record[f\"is_isomorphic_{c_type}\"])))\n",
    "            print()\n",
    "            for k in range(args.topk):\n",
    "                print(\"k_{}:    \\tedit_distance_mean: {:.4f}    acc_iso_mean: {:.4f}\".format(\n",
    "                    k, np.mean(data_record[f\"distance_edit_{k}\"]), np.mean(data_record[f\"is_isomorphic_{k}\"])))\n",
    "            print()\n",
    "            print(\"all:     \\tedit_distance_mean: {:.4f}    acc_iso_mean: {:.4f}\".format(\n",
    "                    np.mean([np.mean(data_record[f\"distance_edit_{c_type}_{k}\"]) for k in range(args.topk) for c_type in concept_labels]),\n",
    "                    np.mean([np.mean(data_record[f\"is_isomorphic_{c_type}_{k}\"]) for k in range(args.topk) for c_type in concept_labels])))\n",
    "            print()\n",
    "            # Printing:\n",
    "            if i % args.inspect_interval == 0 or i == len(dataloader) - 1:\n",
    "                make_dir(dirname + filename)\n",
    "                try_call(pdump, args=[data_record, dirname + filename])\n",
    "                p.print(\"Task {} finished and saved.\\n\".format(i))\n",
    "                print()\n",
    "                sys.stdout.flush()\n",
    "            del data\n",
    "            gc.collect()\n",
    "\n",
    "    if \"classify\" in tasks:\n",
    "        # Obtain the graph for each demonstration for each example:\n",
    "        graphs_gt_all = []\n",
    "        concept_labels_all = []\n",
    "        example_labels_all = []\n",
    "        example_idx_all = []\n",
    "        concept_idx_all = []\n",
    "        for i, data in enumerate(dataloader):\n",
    "            concept_labels = [ele[0] for ele in data[0][2]]\n",
    "            example_labels = [ele[0] for ele in data[1][2]]\n",
    "            example_idx = np.array([concept_labels.index(example_label) for example_label in example_labels])\n",
    "            concept_idx = np.array([example_labels.index(concept_label) for concept_label in concept_labels])\n",
    "\n",
    "            graphs_gt_all.append([get_concept_graph(concept_label, is_new_vertical=args.is_new_vertical, is_concept=args.is_concept, is_relation=args.is_relation) for concept_label in concept_labels])\n",
    "            concept_labels_all.append(concept_labels)\n",
    "            example_labels_all.append(example_labels)\n",
    "            example_idx_all.append(example_idx)\n",
    "            concept_idx_all.append(concept_idx)\n",
    "\n",
    "        # Classify:\n",
    "        Dict = {\"args\": args.__dict__, \"results\": {}}\n",
    "        alpha_list = np.linspace(0,1,21).round(3)\n",
    "        make_dir(dirname + filename)\n",
    "        n_concepts_per_task = len(graphs_gt_all[i])\n",
    "        n_examples_per_task = len(example_labels)\n",
    "\n",
    "        if args.load_parse_src != \"gt\":\n",
    "            parse_result = pload(EXP_PATH + \"/\" + args.load_parse_src)\n",
    "            print(\"Loading parsing result from {}.\".format(EXP_PATH + \"/\" + args.load_parse_src))\n",
    "\n",
    "        args_3D = deepcopy(args)\n",
    "        args_3D.canvas_size = args.canvas_size_3D\n",
    "        args_3D.image_size = (args.canvas_size_3D, args.canvas_size_3D)\n",
    "        for i, data in enumerate(dataloader):\n",
    "            selectors = {}\n",
    "            for j in range(n_concepts_per_task):\n",
    "                if args.load_parse_src == \"gt\":\n",
    "                    graph_j = graphs_gt_all[i][j]\n",
    "                else:\n",
    "                    # Load from file:\n",
    "                    graph_j = parse_result[f'graph_trimmed_ex_{j}_0'][i]\n",
    "                if args.is_bidirectional_re:\n",
    "                    graph_j = get_bidirectional_graph(graph_j)\n",
    "                selectors[j] = get_selector_from_graph(graph_j, ebm_dict_3D, CONCEPTS, OPERATORS)\n",
    "            if i % args.inspect_interval == 0:\n",
    "                p.print(\"iter: {}\".format(i))\n",
    "            energy_dict = {}\n",
    "            input = torch.cat(data[1][0]).to(device)\n",
    "            input_rescaled = F.interpolate(input, size=args.canvas_size_3D, mode=\"nearest\")\n",
    "            mask = torch.cat(data[1][1])\n",
    "            mask_rescaled = F.interpolate(mask, size=args.canvas_size_3D, mode=\"nearest\")\n",
    "            for j in range(n_concepts_per_task):\n",
    "                selector = selectors[j]\n",
    "                for ll in range(n_examples_per_task):\n",
    "                    if isplot:\n",
    "                        p.print(f\"i={i}, j={j}, ll={ll}:\")\n",
    "                    _, masks_top, info = get_selector_SGLD(\n",
    "                        selector, input_rescaled[ll:ll+1], args_3D,\n",
    "                        ebm_target=\"mask\",\n",
    "                        init=args.init,\n",
    "                        SGLD_object_exceed_coef=0,\n",
    "                        SGLD_mutual_exclusive_coef=args.SGLD_mutual_exclusive_coef,\n",
    "                        SGLD_pixel_entropy_coef=args.SGLD_pixel_entropy_coef,\n",
    "                        ensemble_size=args.ensemble_size,\n",
    "                        sample_step=args.sample_step,\n",
    "                        isplot=isplot,\n",
    "                    )  # mask_top: each element [topk, n_examples, 1, H, W]\n",
    "                    energy_dict[(j,ll)] = {}\n",
    "                    energy_dict[(j,ll)][\"masks_top\"] = np.stack(to_np_array(*masks_top, keep_list=True), 2)[0]  # [n_ex, n_masks, 1, H, W]\n",
    "                    energy_dict[(j,ll)][\"masks_max\"] =  energy_dict[(j,ll)][\"masks_top\"].max(1)  # [n_ex, 1, H, W]\n",
    "                    energy_dict[(j,ll)][\"energy_mask\"] = np.stack(list(info['energy_mask'].values()), 2)[0]  # [B, n_masks_concept_re]\n",
    "                    energy_dict[(j,ll)][\"energy\"] = info['energy'][0]\n",
    "                    energy_dict[(j,ll)][\"energy_mean\"] = np.mean(energy_dict[(j,ll)][\"energy_mask\"], -1)  # [B]\n",
    "                    energy_dict[(j,ll)][\"energy_sum\"] = np.sum(energy_dict[(j,ll)][\"energy_mask\"], -1)  # [B]\n",
    "                    energy_dict[(j,ll)][\"mutual_exclusive\"] = info['mutual_exclusive_list'][0]\n",
    "                    for alpha in alpha_list:\n",
    "                        energy_dict[(j,ll)][\"energy_total^m:{:.2f}\".format(alpha)] = energy_dict[(j,ll)][\"energy_mean\"] + energy_dict[(j,ll)][\"mutual_exclusive\"] * alpha\n",
    "                        energy_dict[(j,ll)][\"energy_total^s:{:.2f}\".format(alpha)] = energy_dict[(j,ll)][\"energy_sum\"] + energy_dict[(j,ll)][\"mutual_exclusive\"] * alpha\n",
    "            key_list = [\"masks_top\", \"masks_max\", \"energy\", \"energy_mask\", \"energy_mean\", \"energy_sum\", \"mutual_exclusive\"]\n",
    "            for alpha in alpha_list:\n",
    "                key_list.append(\"energy_total^m:{:.2f}\".format(alpha))\n",
    "                key_list.append(\"energy_total^s:{:.2f}\".format(alpha))\n",
    "            energy_dict_all = {}\n",
    "            for key in key_list:\n",
    "                try:\n",
    "                    energy_dict_all[key] = np.array([np.array([energy_dict[(j,ll)][key] for ll in range(n_examples_per_task)]) for j in range(n_concepts_per_task)])  # [n_concepts, n_ex]\n",
    "                except:\n",
    "                    energy_dict_all[key] = [[energy_dict[(j,ll)][key] for ll in range(n_examples_per_task)] for j in range(n_concepts_per_task)]\n",
    "            concept_idx_pred = energy_dict_all[\"energy\"].squeeze().argmin(1)\n",
    "            concept_idx_gt = concept_idx_all[i]\n",
    "            acc = (concept_idx_pred == concept_idx_gt).mean()\n",
    "            masks_max = energy_dict_all[\"masks_max\"]  # \"masks_max\", : [n_concepts, n_ex, 1, 1, H, W]\n",
    "            mask_chosen = np.concatenate([masks_max[j, concept_idx_pred[j]] for j in range(n_concepts_per_task)])\n",
    "            mask_gt = np.stack([mask_rescaled[concept_idx_pred[j]] for j in range(n_concepts_per_task)])\n",
    "            iou_ex = get_soft_IoU(mask_chosen.round(), mask_gt.round(), dim=(-3,-2,-1))\n",
    "            iou = iou_ex.mean()\n",
    "\n",
    "            # mask_rescaled: [n_ex, 1, H, W]\n",
    "            p.print(\"i={}, acc={:.3f}, iou={:.4f}  pred:{}, gt: {}\".format(i, acc, iou, concept_idx_pred, concept_idx_gt))\n",
    "            record_data(Dict[\"results\"], [acc, concept_idx_pred, concept_idx_gt, iou, iou_ex, mask_rescaled], [\"acc\", \"concept_idx_pred\", \"concept_idx_gt\", \"iou\", \"iou_ex\", \"mask_rescaled\"])\n",
    "            record_data(Dict[\"results\"], list(energy_dict_all.values()), list(energy_dict_all.keys()))\n",
    "            if i % args.inspect_interval == 0 or i == len(dataloader) - 1:\n",
    "                p.print(\"{}: acc={}, iou={}\".format(i, np.mean(Dict[\"results\"][\"acc\"]), np.mean(Dict[\"results\"][\"iou\"])))\n",
    "                try_call(pdump, args=[Dict, dirname + filename], max_exp_time=300)\n",
    "                print()\n",
    "        try_call(pdump, args=[Dict, dirname + filename])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d01576a-cf83-4054-9940-0c459748795f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Random graph:\n",
    "if args.evaluation_type.startswith(\"yc\") and args.model_type == \"rand-graph\":\n",
    "    # Get frequency:\n",
    "    freq_dict = {}\n",
    "    for data in dataset:\n",
    "        c_types = data[0][2]\n",
    "        for c_type in c_types:\n",
    "            graph_gt = get_concept_graph(c_type, is_new_vertical=args.is_new_vertical, is_bidirectional_re=False)\n",
    "            n_line = np.sum([1 for ele in graph_gt if ele[1] == \"Line\"])\n",
    "            n_vertical_mid = np.sum([1 for ele in graph_gt if ele[1] == 'VerticalMid'])\n",
    "            n_vertical_edge = np.sum([1 for ele in graph_gt if ele[1] == 'VerticalEdge'])\n",
    "            n_parallel = np.sum([1 for ele in graph_gt if ele[-1] == 'Parallel'])\n",
    "            record_data(freq_dict, [n_line, n_vertical_mid, n_vertical_edge, n_parallel], [\"Line\", \"VerticalMid\", \"VerticalEdge\", \"Parallel\"])\n",
    "    freq_mean = transform_dict(freq_dict, \"mean\")\n",
    "    freq_vertical_edge = freq_mean[\"VerticalEdge\"]\n",
    "    freq_vertical_mid = freq_mean[\"VerticalMid\"]\n",
    "    freq_parallel = freq_mean['Parallel']\n",
    "    freq_ver_para = np.array([freq_vertical_mid, freq_vertical_edge, freq_parallel])\n",
    "    freq_ver_para = freq_ver_para / freq_ver_para.sum()\n",
    "\n",
    "    # Compute metric:\n",
    "    distance_edit_all = []\n",
    "    is_isomorphic_all = []\n",
    "    for data in dataset:\n",
    "        c_types = data[0][2]\n",
    "        for c_type in c_types:\n",
    "            graph_gt = get_concept_graph(c_type, is_new_vertical=args.is_new_vertical, is_bidirectional_re=False)\n",
    "            graph_rand = get_rand_graph(freq_dict, is_new_vertical=args.is_new_vertical)\n",
    "            distance_edit = get_graph_edit_distance(graph_rand, graph_gt, to_undirected=True)\n",
    "            is_isomorphic = (distance_edit == 0).astype(int)\n",
    "            distance_edit_all.append(distance_edit)\n",
    "            is_isomorphic_all.append(is_isomorphic)\n",
    "    print(\"Edit_distance: {:.4f}\".format(np.mean(distance_edit_all)))\n",
    "    print(\"acc_isomorphic: {:.4f}\".format(np.mean(is_isomorphic_all)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce0eb516-3805-48a9-be92-dfa03dd90747",
   "metadata": {},
   "source": [
    "# 4. Analysis:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b13514c-66ff-411b-9a98-b6250d1ec973",
   "metadata": {},
   "source": [
    "## 4.1 For zero-shot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7d40799-bc51-4a98-9172-70bbd2299d86",
   "metadata": {},
   "outputs": [],
   "source": [
    "isplot = False\n",
    "verbose = 1\n",
    "args.is_analysis = is_jupyter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bd477d8-7721-49e4-958a-227d2b7f37a5",
   "metadata": {},
   "source": [
    "### 4.1.1 Parse:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a1cb76e-1e8d-4770-9b9c-5c7444453f69",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "args.evaluation_type = \"parse\"\n",
    "verbose = True\n",
    "parse_filter_mode = \"mask\"\n",
    "isplot = False\n",
    "show_only_error = True\n",
    "inspect_interval = 500\n",
    "is_baseline = False\n",
    "if args.is_analysis and args.evaluation_type.startswith(\"parse\"):\n",
    "    # # Mask-RCNN, New Vertical, 400 examples:\n",
    "    # graph_list = pload(\"/dfs/user/tailin/.results/BabyARC_baselines/relation_models/mask_rcnn/eval/maskrcnn-conv_d_10-05_e_obj16full-cnn-v3-lr2.5e-5-fix-turing4_m_turing4_Hash_9IeiNpkP.p\")\n",
    "    # is_baseline = False\n",
    "\n",
    "    # # Rand-graph:\n",
    "    # filename = \"evaluation_parse_canvas_16_color_1,2_model_rand-graph_mutu_500_ens_64_sas_150_newv_True_seed_2_id_0_Hash_+hRuVNuO_turing4.p\"\n",
    "    # dirname = \"/dfs/user/tailin/.results/evaluation_parse_12-10/\"\n",
    "    \n",
    "    # Best for new_vertical, Dec 8:\n",
    "    filename = \"evaluation_parse_canvas_16_color_1,2_model_hc-ebm_mutu_500.0_ens_64_sas_150_newv_True_seed_2_id_None_Hash_kEIST2aZ_turing4.p\"\n",
    "    dirname = \"results/evaluation_parse_1-21/\"  # new 400 examples\n",
    "\n",
    "    filenames = filter_filename(dirname, include=[\"canvas_16\"], exclude=[\"Vht+VpvC\"])\n",
    "    # filenames = [filename]\n",
    "    analysis_all = {}\n",
    "    df_dict_list = []\n",
    "    threshold_pixels = 0\n",
    "    thresholds = [None, 0.5, 0.45, 0.4, 0.35, 0.3, 0.25, 0.2, 0.15, 0.1, 0.05, 0]\n",
    "    print(\"Processing {} files.\".format(len(filenames)))\n",
    "    for filename in filenames:\n",
    "        print(filename)\n",
    "        data_record = pload(dirname + filename)\n",
    "        for key in [\"is_new_vertical\", \"concept_model_hash\", \"SGLD_mutual_exclusive_coef\", \"SGLD_pixel_entropy_coef\", \"SGLD_is_anneal\", \"sample_step\", \"ensemble_size\", \"SGLD_is_penalize_lower\", \"is_round_mask\"]:\n",
    "            print(\"{}: {}\".format(key, update_default_hyperparam(data_record[\"args\"])[key]))\n",
    "        analysis_dict = {}\n",
    "\n",
    "        df_dict = {}\n",
    "        df_dict.update(data_record[\"args\"])\n",
    "        df_dict[\"hash\"] = filename.split(\"_\")[-2]\n",
    "        for i in range(len(data_record['parsing_all_dict'][\"graph\"])):\n",
    "            if is_baseline:\n",
    "                data = dataset[i]\n",
    "                c_type_gt = data[3][\"obj_spec\"][0][0][1].split(\"_\")[0]\n",
    "                graph_gt = get_concept_graph(c_type_gt, is_new_vertical=args.is_new_vertical)\n",
    "                graph_pred_ori = graph_list[i]\n",
    "            else:\n",
    "                c_type_gt = data_record['parsing_all_dict'][\"info\"][i][\"obj_spec\"][0][0][1].split(\"_\")[0]\n",
    "                graph_gt = get_concept_graph(c_type_gt, is_new_vertical=data_record[\"args\"][\"is_new_vertical\"])\n",
    "                graph_pred_ori = shorten_graph(data_record['parsing_all_dict'][\"graph\"][i])\n",
    "            mask_top = data_record['parsing_all_dict'][\"masks_top\"][i]\n",
    "            if args.model_type == \"hc-ebm\":\n",
    "                if parse_filter_mode == \"threshold\":\n",
    "                    for threshold in thresholds:\n",
    "                        graph_pred, removed_items = filter_graph_with_threshold(graph_pred_ori, threshold=threshold)\n",
    "                        distance_edit = get_graph_edit_distance(graph_pred, graph_gt, to_undirected=True)\n",
    "                        is_isomorphic = distance_edit == 0\n",
    "                        record_data(analysis_dict, [distance_edit, is_isomorphic], [\"distance_edit:{}\".format(threshold), \"is_isomorphic:{}\".format(threshold)])\n",
    "                elif parse_filter_mode == \"mask\":\n",
    "                    masks_top = data_record['parsing_all_dict'][\"masks_top\"][i]\n",
    "                    masks_is_invalid = [is_mask_invalid(mask, threshold_pixels=threshold_pixels) for mask in masks_top]\n",
    "                    graph_pred = filter_graph_with_masks(graph_pred_ori, masks_is_invalid=masks_is_invalid)\n",
    "                    distance_edit = get_graph_edit_distance(graph_pred, graph_gt, to_undirected=True)\n",
    "                    is_isomorphic = distance_edit == 0\n",
    "                    record_data(analysis_dict, [distance_edit, is_isomorphic], [\"distance_edit\", \"is_isomorphic\"])\n",
    "                else:\n",
    "                    raise\n",
    "            elif args.model_type == \"rand-graph\":\n",
    "                distance_edit = get_graph_edit_distance(graph_pred_ori, graph_gt, to_undirected=True)\n",
    "                is_isomorphic = distance_edit == 0\n",
    "                record_data(analysis_dict, [distance_edit, is_isomorphic], [\"distance_edit\", \"is_isomorphic\"])\n",
    "            if i % inspect_interval == 0 and i > 0 or i == len(data_record['parsing_all_dict'][\"graph\"]) - 1:\n",
    "                print(\"Data {}:\".format(i))\n",
    "                if isplot and (show_only_error is False or is_isomorphic != 1):\n",
    "                    visualize_matrices(data_record['parsing_all_dict'][\"input\"][i].argmax(1))\n",
    "                    plot_matrices(torch.stack(data_record['parsing_all_dict'][\"masks_top\"][i]).squeeze(), images_per_row=6, scale_limit=(0,1))\n",
    "                    print(\"prediction:\")\n",
    "                    pp.pprint(graph_pred)\n",
    "                    print(\"\\nground-truth:\")\n",
    "                    pp.pprint(graph_gt)\n",
    "                if args.model_type == \"hc-ebm\":\n",
    "                    if parse_filter_mode == \"threshold\":\n",
    "                        for threshold in thresholds:\n",
    "                            print(\"Threshold: {}    \\tEdit distance: {}  \\tmean: {:.3f} \\t is_iso: {:.4f}\".format(\n",
    "                                threshold, analysis_dict[\"distance_edit:{}\".format(threshold)][-1],\n",
    "                                np.mean(analysis_dict[\"distance_edit:{}\".format(threshold)]),\n",
    "                                np.mean(analysis_dict[\"is_isomorphic:{}\".format(threshold)])))\n",
    "                    elif parse_filter_mode == \"mask\":\n",
    "                        print(\"Edit distance: {}  \\tmean: {:.3f} \\t is_iso: {:.4f}\".format(\n",
    "                            analysis_dict[\"distance_edit\"][-1],\n",
    "                            np.mean(analysis_dict[\"distance_edit\"]),\n",
    "                            np.mean(analysis_dict[\"is_isomorphic\"])))\n",
    "                elif args.model_type == \"rand-graph\":\n",
    "                    print(\"Edit distance: {}  \\tmean: {:.3f} \\t is_iso: {:.4f}\".format(\n",
    "                        analysis_dict[\"distance_edit\"][-1],\n",
    "                        np.mean(analysis_dict[\"distance_edit\"]),\n",
    "                        np.mean(analysis_dict[\"is_isomorphic\"])))\n",
    "                print(\"\\n\")\n",
    "        record_data(df_dict, [np.mean(analysis_dict[\"distance_edit\"]), np.mean(analysis_dict[\"is_isomorphic\"])], [\"distance_edit\", \"is_isomorphic\"], nolist=True)\n",
    "        df_dict_list.append(df_dict)\n",
    "        analysis_all[filename] = analysis_dict\n",
    "        print(\"\\n\")\n",
    "    analysis_all[\"df\"] = pd.DataFrame(df_dict_list)\n",
    "    df_group = groupby_add_keys(analysis_all[\"df\"], by=[\"is_new_vertical\", \"concept_model_hash\", \"SGLD_mutual_exclusive_coef\", \"SGLD_pixel_entropy_coef\", \"SGLD_is_anneal\", \"SGLD_is_penalize_lower\", \"sample_step\", \"ensemble_size\", \"is_round_mask\"],\n",
    "        add_keys=\"hash\",\n",
    "        other_keys=[\"is_isomorphic\", \"distance_edit\"],\n",
    "        mode={\"max\": [\"is_isomorphic\"], \"min\":[\"distance_edit\"], \"count\": None}\n",
    "    )\n",
    "    display(df_group.style.background_gradient(cmap=plt.cm.get_cmap(\"PiYG\")))\n",
    "    pdump(analysis_all, dirname + \"analysis_all.p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef769dbd-f25f-4de0-af8c-7490bef882a0",
   "metadata": {},
   "source": [
    "### 4.1.2 Grounding:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19bbcae3-b556-4460-add7-07e950ac564c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "args.evaluation_type = \"grounding\"\n",
    "total_tasks = 200\n",
    "if args.is_analysis and args.evaluation_type.startswith(\"grounding\"):\n",
    "    df_dict_list = []\n",
    "    for c_type in [\"Eshape\", \"Fshape\", \"Ashape\", \"RectE1a\", \"RectE2a\", \"RectE3a\"]:\n",
    "        dirname = EXP_PATH + \"/evaluation_grounding-{}_1-21/\".format(c_type)\n",
    "        try:\n",
    "            filenames = filter_filename(dirname, include=\".p\")\n",
    "        except:\n",
    "            continue\n",
    "        for filename in filenames:\n",
    "            try:\n",
    "                data_record = pload(dirname + filename)\n",
    "            except Exception as e:\n",
    "                print(f\"Error {e} happens for file {filename}\")\n",
    "            iou = np.mean(data_record[\"iou\"][:total_tasks])\n",
    "            df_dict = update_default_hyperparam_generalization(data_record[\"args\"])\n",
    "            df_dict[\"hash\"] = filename.split(\"_\")[-2]\n",
    "            df_dict[\"machine\"] = filename.split(\"_\")[-1][:-2]\n",
    "            df_dict[\"iou\"] = iou\n",
    "            df_dict[\"epoch\"] = len(data_record[\"iou\"][:total_tasks])\n",
    "            df_dict[\"load_epoch\"] = data_record[\"args\"][\"relation_load_id\"] * 5\n",
    "            df_dict_list.append(df_dict)\n",
    "\n",
    "    df = pd.DataFrame(df_dict_list)\n",
    "    df_group = groupby_add_keys(\n",
    "        df,\n",
    "        by=[\"evaluation_type\",\n",
    "            # 'SGLD_mutual_exclusive_coef',\n",
    "            # 'SGLD_pixel_entropy_coef',\n",
    "            # 'ensemble_size',\n",
    "            # 'sample_step',\n",
    "            # \"SGLD_is_penalize_lower\",\n",
    "            \"concept_model_hash\",\n",
    "            \"relation_model_hash\",\n",
    "            # \"is_bidirectional_re\",\n",
    "            # \"max_n_distractors\",\n",
    "            # \"min_n_distractors\",\n",
    "            \"allow_connect\",\n",
    "            \"is_proper_size\",\n",
    "            \"is_concept\",\n",
    "            \"is_relation\",\n",
    "            # \"id\",\n",
    "            \"gpuid\",\n",
    "           ],\n",
    "        add_keys=[\"hash\", \"machine\"], other_keys=[\"iou\", \"epoch\", \"load_epoch\"],\n",
    "        mode={\n",
    "            \"max\": [\"iou\", \"epoch\", \"load_epoch\"],\n",
    "            \"count\": None,\n",
    "        })\n",
    "    display(df_group.style.background_gradient(cmap=plt.cm.get_cmap(\"PiYG\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7966110d-d907-48e7-8ca9-66270aaacc7b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if args.is_analysis and args.evaluation_type.startswith(\"grounding\"):\n",
    "    dirname = EXP_PATH + \"/evaluation_grounding-{}_1-21/\".format(\"Eshape\")\n",
    "    hash_str = \"1AMKX9yn\"\n",
    "    filename = filter_filename(dirname, include=hash_str)[0]\n",
    "    data_record = pload(dirname + \"/\" + filename)\n",
    "    for i, iou in enumerate(data_record[\"iou\"]):\n",
    "        visualize_matrices(data_record[\"input\"][i].argmax(1), subtitles=[\"{}, iou: {}\".format(i, iou)], images_per_row=6)\n",
    "        plot_matrices(np.concatenate([data_record[\"masks_top\"][i].sum(1).squeeze(0),\n",
    "                    data_record[\"masks_top\"][i].squeeze((0,2))]), images_per_row=6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "86495a7e-014c-4418-9cf5-65df9c73c402",
   "metadata": {},
   "source": [
    "### 4.1.3 Classify:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8856c65b-d635-4db1-9b27-1fbbc8240274",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "args.evaluation_type = \"classify\"\n",
    "verbose = 0\n",
    "total_tasks = 200\n",
    "if args.is_analysis and args.evaluation_type.startswith(\"classify\"):\n",
    "    dirname = EXP_PATH + \"/evaluation_classify_1-21/\"\n",
    "    filenames = filter_filename(dirname, include=\".p\", exclude=\"df_\")\n",
    "    df_dict_list = []\n",
    "    for filename in filenames:\n",
    "        df_dict = {}\n",
    "        if verbose > 0:\n",
    "            p.print(\"{}:\".format(filename), banner_size=100)\n",
    "        data_record = pload(dirname + filename)\n",
    "        if \"acc^s:0.00\" not in data_record[\"results\"]:\n",
    "            if verbose > 0:\n",
    "                print(\"{} not fully processed.\".format(filename))\n",
    "            continue\n",
    "        df_dict.update(update_default_hyperparam_generalization(data_record[\"args\"]))\n",
    "        df_dict[\"load_epoch\"] = data_record[\"args\"][\"relation_load_id\"] * 5\n",
    "        df_dict[\"filename\"] = filename\n",
    "        df_dict[\"hash\"] = filename.split(\"_\")[-2]\n",
    "        df_dict[\"machine\"] = filename.split(\"_\")[-1][:-2]\n",
    "        if verbose > 0:\n",
    "            for key in [\"SGLD_mutual_exclusive_coef\", \"SGLD_is_penalize_lower\", \"concept_model_hash\", \"relation_model_hash\"]:\n",
    "                print(\"{}: {}\".format(key, data_record[\"args\"][key]))\n",
    "        acc_s_list = []\n",
    "        alpha_s_list = []\n",
    "        df_dict[\"n_examples\"] = len(data_record[\"results\"][\"acc^s:0.00\"][:total_tasks])\n",
    "        for alpha in np.linspace(0,1,11):\n",
    "            if verbose > 0:\n",
    "                print(\"acc^s:{:.2f}: {:.3f}\".format(alpha, np.mean(data_record[\"results\"][\"acc^s:{:.2f}\".format(alpha)][:total_tasks])))\n",
    "            acc_s_list.append(np.mean(data_record[\"results\"][\"acc^s:{:.2f}\".format(alpha)][:total_tasks]))\n",
    "            alpha_s_list.append(alpha)\n",
    "        df_dict[\"acc_s_max\"] = np.max(acc_s_list)\n",
    "        df_dict[\"acc_s_argmax\"] = alpha_s_list[np.argmax(acc_s_list)]\n",
    "        df_dict[\"acc_s_0\"] = acc_s_list[0]\n",
    "        df_dict[\"acc_s_0.1\"] = acc_s_list[1]\n",
    "        if verbose > 0:\n",
    "            print()\n",
    "        acc_m_list = []\n",
    "        alpha_m_list = []\n",
    "        for alpha in np.linspace(0,1,11):\n",
    "            if verbose > 0:\n",
    "                print(\"acc^m:{:.2f}: {:.3f}\".format(alpha, np.mean(data_record[\"results\"][\"acc^m:{:.2f}\".format(alpha)][:total_tasks])))\n",
    "            acc_m_list.append(np.mean(data_record[\"results\"][\"acc^m:{:.2f}\".format(alpha)][:total_tasks]))\n",
    "            alpha_m_list.append(alpha)\n",
    "        df_dict[\"acc_m_max\"] = np.max(acc_m_list)\n",
    "        df_dict[\"acc_m_argmax\"] = alpha_m_list[np.argmax(acc_m_list)]\n",
    "        df_dict[\"acc_m_0\"] = acc_m_list[0]\n",
    "        df_dict[\"acc_m_0.1\"] = acc_m_list[1]\n",
    "        df_dict_list.append(df_dict)\n",
    "    df = pd.DataFrame(df_dict_list)\n",
    "    df_group = groupby_add_keys(\n",
    "        df,\n",
    "        by=[\"dataset\",\n",
    "            # \"SGLD_mutual_exclusive_coef\",\n",
    "            # \"SGLD_is_penalize_lower\",\n",
    "            \"concept_model_hash\",\n",
    "            \"relation_model_hash\",\n",
    "            # \"ensemble_size\",\n",
    "            # \"is_bidirectional_re\",\n",
    "            # \"max_n_distractors\",\n",
    "            # \"id\",\n",
    "            \"gpuid\",\n",
    "           ],\n",
    "        add_keys=[\"hash\", \"machine\"],\n",
    "        other_keys=[ \"acc_s_0\", \"acc_s_0.1\", \"acc_s_argmax\", \"acc_m_max\", \"acc_m_argmax\", \"n_examples\", \"load_epoch\"],\n",
    "        mode={\n",
    "            \"mean\": [ \"acc_s_0\", \"acc_s_0.1\", \"acc_s_argmax\", \"acc_m_argmax\", \"n_examples\", \"load_epoch\"],\n",
    "            \"max\": [\"acc_s_max\", \"acc_m_max\"],\n",
    "            \"count\": None,\n",
    "        }\n",
    "    )\n",
    "    pdump({\"df\": df, \"df_group\": df_group}, dirname + \"df_{}-{}.p\".format(datetime.now().month, datetime.now().day))\n",
    "    display(df_group.style.background_gradient(cmap=plt.cm.get_cmap(\"PiYG\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8398de95-2dfa-4136-bdab-0a8b93b1ca56",
   "metadata": {},
   "source": [
    "### 4.2.1 Parse2D+classify3D:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a642637c-8053-4018-8407-7c58f9a227fd",
   "metadata": {},
   "source": [
    "#### 4.2.1.1 Parse:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "962a87ae-0068-4b8f-884b-7f5e699757b6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## 1-1:\n",
    "args.evaluation_type = \"yc-parse+classify^parse\"\n",
    "verbose = 1\n",
    "energy_criterion = {\n",
    "    \"energy\": 2,\n",
    "    \"pixels_under\": 2,\n",
    "    \"pixels_exceed_in\": 2,\n",
    "    \"pixels_exceed_out_mean\": 2,\n",
    "    \"pairwise_distance_topk_mean\": 0.1,\n",
    "}\n",
    "isplot = False\n",
    "def get_instance_energy(df_dict, energy_criterion):\n",
    "    energy = 0\n",
    "    for key, value in energy_criterion.items():\n",
    "        energy += df_dict[key] * value\n",
    "    return energy\n",
    "dff_dict_list = []\n",
    "args.is_analysis = is_jupyter\n",
    "if args.is_analysis and args.evaluation_type.startswith(\"yc-parse+classify\"):\n",
    "    dirname = EXP_PATH + \"/evaluation_yc-parse+classify^parse_1-21/\"\n",
    "    filenames = filter_filename(dirname, include=\".p\", exclude=\"df_\")\n",
    "    # filenames = [\"evaluation_yc-parse+classify^parse_canvas_16_color_1,2_ex_400_min_0_model_hc-ebm_mutu_500.0_ens_64_sas_150_newv_True_batch_1_con_fRZtzn33_re_Wfxw19nM_bi_True_seed_2_id_None_Hash_mU7ILNWm_turing3.p\"]\n",
    "    # filenames = [\"evaluation_yc-parse+classify^parse_canvas_16_color_1,2_ex_400_min_0_model_hc-ebm_mutu_500.0_ens_64_sas_150_newv_True_batch_1_con_4Qb0Vu0x_re_Wfxw19nM_bi_True_seed_2_id_None_Hash_84sNKLJz_hyperturing1.p\"]  # Without L^em\n",
    "    for filename in filenames:\n",
    "        dff_dict = {}\n",
    "        p.print(filename, banner_size=100, is_datetime=False)\n",
    "        try:\n",
    "            data_record = pload(dirname + filename)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            continue\n",
    "        if 'is_isomorphic_0' not in data_record:\n",
    "            print(\"'is_isomorphic_0' not in data_record, skip.\")\n",
    "            continue\n",
    "        args = init_args(update_default_hyperparam_generalization(update_default_hyperparam(data_record[\"args\"])))\n",
    "        c_types = get_c_core(args.dataset.split(\"-\")[1].split(\"+\"))\n",
    "        print(f\"concept_hash: {args.concept_model_hash}\")\n",
    "        print(f\"relation_hash: {args.relation_model_hash}\")\n",
    "        print(f\"SGLD_is_penalize_lower: {args.SGLD_is_penalize_lower}\")\n",
    "        print(f\"init: {args.init}\")\n",
    "        print(f\"infer_order: {args.infer_order}\")\n",
    "        print(f\"SGLD_is_penalize_lower_seq: {args.SGLD_is_penalize_lower_seq}\")\n",
    "        print(f\"lambd: {args.lambd}\")\n",
    "        print(f\"is_bidirectional_re: {args.is_bidirectional_re}\")\n",
    "        print(f\"tasks completed: {len(data_record[f'graph_trimmed_{c_types[0]}_0'])}\")\n",
    "        x = np.arange(args.topk)\n",
    "        acc_rank = [np.mean(data_record[f\"is_isomorphic_{i}\"]) for i in range(args.topk)]\n",
    "        distance_rank = [np.mean(data_record[f\"distance_edit_{i}\"]) for i in range(args.topk)]\n",
    "        # plot_simple(y=acc_rank, title=\"acc vs. rank\", xlabel=\"rank\", ylabel=\"acc\")\n",
    "        plot_2_axis(x=x, y1=acc_rank, y2=distance_rank,\n",
    "                    xlabel=\"rank\", ylabel1=\"acc\", ylabel2=\"distance\",\n",
    "                    ylim1=(0,0.8), ylim2=(0,2.5),\n",
    "                    title=\"acc & distance vs. rank\",\n",
    "                   )\n",
    "\n",
    "        \"\"\"\n",
    "        For each time step, compute the acc, and visualize:\n",
    "        \"\"\"\n",
    "        Dict = {}\n",
    "        Dict_all = {}\n",
    "        best_acc_dict = {}\n",
    "        argmax_dict = {}\n",
    "        graph_distance_min_acc_dict = {}\n",
    "        graph_distance_mode_acc_dict = {}\n",
    "        energy_dict = {}\n",
    "        concept_energy_dict = {}\n",
    "        energy_min_acc_dict = {}\n",
    "        concept_energy_min_acc_dict = {}\n",
    "        df_dict_list = []\n",
    "        for j, c_type in enumerate(c_types):\n",
    "            Dict[c_type] = []\n",
    "            Dict_all[c_type] = []\n",
    "            argmax_dict[c_type] = []\n",
    "            energy_dict[c_type] = []\n",
    "            concept_energy_dict[c_type] = []\n",
    "            if isplot:\n",
    "                print(f\"concept: {c_type}\")\n",
    "            for i in range(len(data_record['distance_edit_0'])):\n",
    "                List_i = []\n",
    "                energy_i = []\n",
    "                concept_energy_i = []\n",
    "                instance_energy_list = []\n",
    "                for k in range(args.topk):\n",
    "                    List_i.append(float(data_record[f\"is_isomorphic_{c_type}_{k}\"][i]))\n",
    "                    graph_trimmed = data_record[f\"graph_trimmed_{c_type}_{k}\"][i]\n",
    "                    energy_i.append(np.sum([ele[2] for ele in graph_trimmed]))\n",
    "                    concept_energy_i.append(np.sum([ele[2] for ele in graph_trimmed if not isinstance(ele[0], tuple)]))\n",
    "\n",
    "                    # Conctruct df_dict:\n",
    "                    df_dict = {}\n",
    "                    df_dict[\"task\"] = i\n",
    "                    df_dict[\"c_type\"] = c_type\n",
    "                    df_dict[\"energy\"] = data_record[f\"energy_{c_type}\"][i][k]\n",
    "                    df_dict[\"rank\"] = k\n",
    "                    df_dict[\"concept_energy\"] = data_record[f\"concept_energy_{c_type}\"][i][k]\n",
    "                    df_dict[\"relation_energy\"] = data_record[f\"relation_energy_{c_type}\"][i][k]\n",
    "                    df_dict[\"distance_edit\"] = data_record[f\"distance_edit_{c_type}_{k}\"][i]\n",
    "                    df_dict[\"is_isomorphic\"] = data_record[f\"is_isomorphic_{c_type}_{k}\"][i]\n",
    "                    df_dict[\"pixels_under\"] = data_record[f\"pixels_under_{c_type}\"][i][k]\n",
    "                    df_dict[\"pixels_exceed_in\"] = data_record[f\"pixels_exceed_in_{c_type}\"][i][k]\n",
    "                    df_dict[\"pixels_exceed_out\"] = data_record[f\"pixels_exceed_out_{c_type}\"][i][k]\n",
    "                    df_dict[\"pixels_exceed_out_mean\"] = data_record[f\"pixels_exceed_out_mean_{c_type}\"][i][k]\n",
    "                    df_dict[\"pairwise_distance_topk\"] = data_record[f\"pairwise_distance_topk_{c_type}\"][i][k]\n",
    "                    df_dict[\"pairwise_distance_topk_mean\"] = data_record[f\"pairwise_distance_topk_{c_type}\"][i][k].mean(-1)\n",
    "                    df_dict[\"graph_trimmed\"] = data_record[f\"graph_trimmed_{c_type}_{k}\"][i]\n",
    "                    df_dict[\"graph_gt\"] = data_record[f\"graph_gt_{c_type}\"][i]\n",
    "                    instance_energy_list.append(get_instance_energy(df_dict, energy_criterion))\n",
    "                    df_dict_list.append(df_dict)\n",
    "\n",
    "                    # Plot the top instance:\n",
    "                    if isplot and k == 0 and df_dict[\"is_isomorphic\"] != True:\n",
    "                        p.print(f\"Task {i}: graph_trimmed, distance_edit = {df_dict['distance_edit']}\", banner_size=60, is_datetime=False)\n",
    "                        p.print(df_dict[\"graph_trimmed\"], is_datetime=False)\n",
    "                        plot_matrices(data_record[f\"gt_mask_{c_type}\"][i].squeeze(1), images_per_row=8)\n",
    "                        plot_matrices(data_record[f\"masks_top_{c_type}\"][i][:, k].squeeze(1), images_per_row=8)\n",
    "                instance_energy_list = np.array(instance_energy_list)\n",
    "                instance_energy_argsort = instance_energy_list.argsort()\n",
    "                for kk, idx in enumerate(instance_energy_argsort):\n",
    "                    df_dict_list[idx - args.topk][\"instance_energy_rank\"] = kk\n",
    "\n",
    "                Dict[c_type].append(max(List_i))\n",
    "                Dict_all[c_type].append(List_i)\n",
    "                argmax_dict[c_type].append(np.argmax(List_i))\n",
    "                energy_dict[c_type].append(energy_i)\n",
    "                concept_energy_dict[c_type].append(concept_energy_i)\n",
    "            best_acc_dict[c_type] = np.mean(Dict[c_type])\n",
    "            Dict_all[c_type] = np.array(Dict_all[c_type])\n",
    "            energy_dict[c_type] = np.array(energy_dict[c_type])\n",
    "            concept_energy_dict[c_type] = np.array(concept_energy_dict[c_type])\n",
    "            graph_distance_min_id = np.array(data_record[f\"pairwise_distance_topk_{c_types[0]}\"]).mean(-1).argmin(-1)\n",
    "            array = np.array(data_record[f\"pairwise_distance_topk_{c_types[0]}\"]).mean(-1)\n",
    "            mode = stats.mode(array, axis=1)[0]\n",
    "            graph_distance_mode_id = np.argmax(mode==array, axis=1)\n",
    "            graph_distance_min_acc_dict[c_type] = np.mean([Dict_all[c_type][i, idx] for i, idx in enumerate(graph_distance_min_id)])\n",
    "            graph_distance_mode_acc_dict[c_type] = np.mean([Dict_all[c_type][i, idx] for i, idx in enumerate(graph_distance_mode_id)])\n",
    "            energy_min_acc_dict[c_type] = np.mean([Dict_all[c_type][i, idx] for i, idx in enumerate(energy_dict[c_type].argmin(1))])\n",
    "            concept_energy_min_acc_dict[c_type] = np.mean([Dict_all[c_type][i, idx] for i, idx in enumerate(concept_energy_dict[c_type].argmin(1))])\n",
    "\n",
    "        df = pd.DataFrame(df_dict_list)\n",
    "        rank1_acc_dict = {c_type: np.mean(data_record[f\"is_isomorphic_{c_type}_{0}\"]) for c_type in c_types}\n",
    "        rank1_distance_edit_dict = {c_type: np.mean(data_record[f\"distance_edit_{c_type}_{0}\"]) for c_type in c_types}\n",
    "        for key in [\"concept_model_hash\", \"relation_model_hash\", \"SGLD_is_penalize_lower\", \"SGLD_is_penalize_lower_seq\",\n",
    "                   \"init\", \"infer_order\", \"lambd\", \"is_bidirectional_re\"]:\n",
    "            dff_dict[key] = getattr(args, key)\n",
    "        dff_dict[\"tasks_completed\"] = len(data_record[f'graph_trimmed_{c_types[0]}_0'])\n",
    "        dff_dict[\"hash\"] = filename.split(\"_\")[-2]\n",
    "\n",
    "        dff_dict[\"acc_rank1\"] = np.mean(list(rank1_acc_dict.values()))\n",
    "        dff_dict[\"distance_edit_rank1\"] = np.mean(list(rank1_distance_edit_dict.values()))\n",
    "        dff_dict[\"acc_weighted\"] = df.groupby(by=[\"instance_energy_rank\"]).mean().iloc[0][\"is_isomorphic\"]\n",
    "        dff_dict[\"acc_best\"] = np.mean(list(best_acc_dict.values()))\n",
    "        dff_dict[\"acc_graph_min\"] = np.mean(list(graph_distance_min_acc_dict.values()))\n",
    "        dff_dict[\"acc_graph_mode\"] = np.mean(list(graph_distance_mode_acc_dict.values()))\n",
    "        dff_dict[\"acc_energy_min\"] = np.mean(list(energy_min_acc_dict.values()))\n",
    "        dff_dict[\"acc_concept_energy_min\"] = np.mean(list(concept_energy_min_acc_dict.values()))\n",
    "        dff_dict_list.append(dff_dict)\n",
    "        display(df.groupby(by=\"c_type\").mean())\n",
    "        if verbose >= 2:\n",
    "            print(\"Acc for rank-1 model:\")\n",
    "            pp.pprint(rank1_acc_dict)\n",
    "            print(\"Avg rank-1 acc: {:.4f}\".format(np.mean(list(rank1_acc_dict.values()))))\n",
    "            print()\n",
    "            print(\"Acc for graph_distance_min:\")\n",
    "            pp.pprint(graph_distance_min_acc_dict)\n",
    "            print(\"Avg graph_distance_min acc: {:.4f}\".format(np.mean(list(graph_distance_min_acc_dict.values()))))\n",
    "            print()\n",
    "            print(\"Acc for graph_distance_mode:\")\n",
    "            pp.pprint(graph_distance_mode_acc_dict)\n",
    "            print(\"Avg graph_distance_mode acc: {:.4f}\".format(np.mean(list(graph_distance_mode_acc_dict.values()))))\n",
    "            print()\n",
    "            print(\"Acc for energy_min:\")\n",
    "            pp.pprint(energy_min_acc_dict)\n",
    "            print(\"Avg energy_min acc: {:.4f}\".format(np.mean(list(energy_min_acc_dict.values()))))\n",
    "            print()\n",
    "            print(\"Acc for concept_energy_min:\")\n",
    "            pp.pprint(concept_energy_min_acc_dict)\n",
    "            print(\"Avg concept_energy_min acc: {:.4f}\".format(np.mean(list(concept_energy_min_acc_dict.values()))))\n",
    "            print()\n",
    "            print(\"Best acc:\")\n",
    "            pp.pprint(best_acc_dict)\n",
    "            print(\"Avg best acc: {:.4f}\".format(np.mean(list(best_acc_dict.values()))))\n",
    "            print()\n",
    "\n",
    "            display(df.groupby(by=[\"instance_energy_rank\"]).mean().style.background_gradient(cmap=plt.cm.get_cmap(\"PiYG_r\")))\n",
    "    dff = pd.DataFrame(dff_dict_list)\n",
    "    mean_edit_distance = np.mean([\n",
    "        np.mean(data_record[\"distance_edit_ex_0_0\"]),\n",
    "        np.mean(data_record[\"distance_edit_ex_1_0\"]),\n",
    "        np.mean(data_record[\"distance_edit_ex_2_0\"])])\n",
    "    acc_top1 = np.mean(data_record[f\"is_isomorphic_0\"])\n",
    "    print(\"iso_acc: {:.4f}    edit_distance: {:.4f}\".format(acc_top1, mean_edit_distance))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9057d5d4-46f0-4891-87b4-7b51926b7f3a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if args.is_analysis and args.evaluation_type.startswith(\"yc-parse+classify\"):\n",
    "    keys = [\"concept_model_hash\", \"relation_model_hash\"]\n",
    "    display(dff.groupby(by=keys).mean().style.background_gradient(cmap=plt.cm.get_cmap(\"PiYG\")))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "280491fd-0ec4-46e7-9023-273cbc908110",
   "metadata": {},
   "source": [
    "#### 4.2.1.2 Classify:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7f99af1-0622-4a70-86b9-ea9863d0c579",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if args.is_analysis and args.evaluation_type.startswith(\"yc-parse+classify\"):\n",
    "    dirname = EXP_PATH + \"/evaluation_yc-parse+classify^classify_1-21/\"\n",
    "    filenames = filter_filename(dirname, include=\".p\", exclude=\"df_\")\n",
    "    df_dict_list = []\n",
    "    for filename in filenames:\n",
    "        df_dict = {}\n",
    "        p.print(filename, is_datetime=False)\n",
    "        try:\n",
    "            data_record = pload(dirname + filename)\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            continue\n",
    "        df_dict = deepcopy(update_default_hyperparam_generalization(data_record[\"args\"]))\n",
    "        df_dict[\"acc\"] = np.mean(data_record[\"results\"][\"acc\"])\n",
    "        df_dict[\"iou\"] = np.mean(data_record[\"results\"][\"iou\"])\n",
    "        df_dict[\"tasks\"] = len(data_record[\"results\"][\"iou\"])\n",
    "        # df_dict[\"hash\"] = \n",
    "        df_dict_list.append(df_dict)\n",
    "    df = pd.DataFrame(df_dict_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae559712-5c17-42f0-93f6-4674de55031a",
   "metadata": {},
   "outputs": [],
   "source": [
    "if args.is_analysis and args.evaluation_type.startswith(\"yc-parse+classify\"):\n",
    "    df_group = groupby_add_keys(\n",
    "            df,\n",
    "            by=[\"load_parse_src\",\"ensemble_size\", \"gpuid\",\n",
    "               ],\n",
    "            add_keys=[],\n",
    "            other_keys=[\"acc\", \"iou\",\"tasks\",],\n",
    "            mode={\n",
    "                \"mean\": [\"acc\", \"iou\",\"tasks\"],\n",
    "                \"count\": None,\n",
    "            }\n",
    "        )\n",
    "    pdump({\"df\": df, \"df_group\": df_group}, dirname + \"df_{}-{}.p\".format(datetime.now().month, datetime.now().day))\n",
    "    display(df_group.style.background_gradient(cmap=plt.cm.get_cmap(\"PiYG\")))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "base"
  },
  "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,
  "toc-showmarkdowntxt": true,
  "toc-showtags": false
 },
 "nbformat": 4,
 "nbformat_minor": 5
}