principle_2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "import json\n",
    "import openai\n",
    "import random\n",
    "import os\n",
    "import collections\n",
    "import copy\n",
    "import netgraph\n",
    "import seaborn as sns\n",
    "import scipy.stats\n",
    "import replicate\n",
    "from utils import get_response, summarize_reasons\n",
    "\n",
    "MEDIUM_SIZE = 26\n",
    "SMALL_SIZE = 0.85 * MEDIUM_SIZE\n",
    "BIGGER_SIZE = 1.5 * MEDIUM_SIZE\n",
    "\n",
    "plt.rc('font', size=SMALL_SIZE)          # controls default text sizes\n",
    "plt.rc('axes', titlesize=SMALL_SIZE)     # fontsize of the axes title\n",
    "plt.rc('axes', labelsize=MEDIUM_SIZE)    # fontsize of the x and y labels\n",
    "plt.rc('xtick', labelsize=0.7*SMALL_SIZE)    # fontsize of the tick labels\n",
    "plt.rc('ytick', labelsize=0.7*SMALL_SIZE)    # fontsize of the tick labels\n",
    "plt.rc('legend', fontsize=SMALL_SIZE)    # legend fontsize\n",
    "plt.rc('figure', titlesize=BIGGER_SIZE)  # fontsize of the figure title\n",
    "\n",
    "def network_growth(n0, temperature, model, environment, role, method='llm', num_common_neighbors=True, cot=False):\n",
    "\n",
    "    if method == 'sbm':\n",
    "        G = nx.stochastic_block_model([n0 // 2, n0 // 2], [[0.5, 0.1], [0.1, 0.5]])\n",
    "   \n",
    "        return [G], []  \n",
    "    else:\n",
    "        G = nx.stochastic_block_model([n0 // 2, n0 // 2], [[0.5, 0.1], [0.1, 0.5]], seed=0)\n",
    "\n",
    "\n",
    "    Gs = []\n",
    "    results = []\n",
    "\n",
    "    for t in G.nodes():\n",
    "\n",
    "        if method == 'llm':\n",
    "            result = select_neighbor(G, t, temperature, num_common_neighbors=num_common_neighbors, model=model, environment=environment, role=role, cot=cot)\n",
    "            if result:\n",
    "                v = result['name']\n",
    "                G.add_edge(t, v)\n",
    "                results.append(result)\n",
    "        elif method == 'random':\n",
    "            v = random.choice(list(set(G.nodes() - set(G.neighbors(t)))))\n",
    "            G.add_edge(t, v)\n",
    "            results.append({'name' : v, 'common_friends' : list(set(G.neighbors(v)) & set(G.neighbors(t))), 'reason' : 'random'})\n",
    "        elif method == 'winner':\n",
    "            v = None,\n",
    "            max_common_friends = 0\n",
    "            for u in G.nodes():\n",
    "                if u not in G.neighbors(t) and u != t and len(set(G.neighbors(u)) & set(G.neighbors(t))) > max_common_friends:\n",
    "                    v = u\n",
    "                    max_common_friends = len(set(G.neighbors(u)) & set(G.neighbors(t)))\n",
    "\n",
    "            G.add_edge(t, v)\n",
    "            results.append({'name' : v, 'common_friends' : list(set(G.neighbors(v)) & set(G.neighbors(t))), 'reason' : 'winner'})\n",
    "\n",
    "        Gs.append(G.copy())\n",
    "\n",
    "    return Gs, results\n",
    "\n",
    "def select_neighbor(G, t, temperature, model, environment, role, num_common_neighbors, cot):\n",
    "    \n",
    "    candidate_profiles = []\n",
    "    for v in G.nodes():\n",
    "        if v != t and v not in G.neighbors(t):\n",
    "            if num_common_neighbors:\n",
    "                candidate_profiles.append({'name' : v, 'common_friends' : len(set(G.neighbors(v)) & set(G.neighbors(t)))})\n",
    "            else:\n",
    "                candidate_profiles.append({'name' : v, 'friends' : list(G.neighbors(v))})\n",
    "\n",
    "    if cot: \n",
    "        output_format = f\"\"\"\n",
    "    {{\n",
    "        \"reason\" : reason for selecting the person,\n",
    "        \"name\" : name of the person you selected\n",
    "    }}\n",
    "        \"\"\"\n",
    "    else:\n",
    "        output_format = f\"\"\"\n",
    "    {{\n",
    "        \"name\" : name of the person you selected,\n",
    "        \"reason\" : reason for selecting the person\n",
    "    }}\n",
    "        \"\"\"\n",
    "\n",
    "    prompt = f\"\"\"\n",
    "    # Task\n",
    "    {f'You are in a {environment}.' if environment else ''}Your task is to select a person to be {role} with.\n",
    "\n",
    "    # Input\n",
    "    The input is a list of dictionaries. \n",
    "    Your profile is given below after chevrons:\n",
    "    <PROFILE>\n",
    "    {json.dumps({'name' : t, 'friends' : list(G.neighbors(t))})}\n",
    "    </PROFILE>\n",
    "\n",
    "    The list of candidate profiles is given below after chevrons:\n",
    "    <PROFILES>\n",
    "    {json.dumps(candidate_profiles)}\n",
    "    </PROFILES>\n",
    "\n",
    "    # Output\n",
    "    The output should be given in JSON format with the following structure\n",
    "\n",
    "    {output_format}\n",
    "\n",
    "    # Notes\n",
    "    * The name of the person you selected must be one of the names in the input.\n",
    "    * Your output must be JSON only.\n",
    "\n",
    "    ```json\n",
    "    \"\"\"   \n",
    "\n",
    "    for i in range(10):\n",
    "        try:\n",
    "            ans = get_response(prompt, temperature=temperature, system_prompt=\"You are a helpful assistant\", model=model)\n",
    "\n",
    "            try:\n",
    "                result = json.loads(ans.split('```')[0])\n",
    "            except:\n",
    "                result = json.loads(ans.split('```json')[1].split('```')[0])\n",
    "\n",
    "            if result['name'] in G.nodes():    \n",
    "                print('NEW EDGE', result)\n",
    "                return result \n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "    \n",
    "def run_network_formation_experiment(n_min, n_max, n_step, num_simulations, outfile, temperatures, method, model, environment=None, role='friends', num_common_neighbors=True, cot=False):\n",
    "    saved_scenarios = set()\n",
    "\n",
    "    if os.path.exists(outfile):\n",
    "        with open(outfile) as f:\n",
    "            lines = f.read().splitlines()\n",
    "\n",
    "            for line in lines:\n",
    "                scenario = json.loads(line)\n",
    "                saved_scenarios.add((scenario['n'], scenario['simulation'], scenario['temperature']))\n",
    "\n",
    "        exit()\n",
    "\n",
    "    f = open(outfile, 'a+')\n",
    "\n",
    "\n",
    "    for n in range(n_min, n_max + 1, n_step):\n",
    "        for i in range(num_simulations):\n",
    "            for temperature in temperatures:\n",
    "                if (n, i, temperature) in saved_scenarios:\n",
    "                    print(f'Skipping simulation for n={n}, i={i}, temperature={temperature}')\n",
    "                    continue\n",
    "                else:\n",
    "                    print(f'Running simulation for n={n}, i={i}, temperature={temperature}')\n",
    "\n",
    "                    Gs, reasons = network_growth(n, temperature=temperature, method=method, num_common_neighbors=num_common_neighbors, model=model, environment=environment, role=role, cot=cot)\n",
    "\n",
    "                    temp = {\n",
    "                        'n' : n,\n",
    "                        'temperature' : temperature,\n",
    "                        'simulation' : i,\n",
    "                        'graphs' : [nx.to_dict_of_lists(G) for G in Gs],\n",
    "                        'reasons' : reasons\n",
    "                    }    \n",
    "\n",
    "                    f.write(json.dumps(temp) + '\\n')            \n",
    "\n",
    "                if method != 'llm':\n",
    "                    break\n",
    "\n",
    "    f.close()\n",
    "\n",
    "def draw_graph(G, ax, G0=None, use_netgraph=True):\n",
    "    group_1 = [n for n in G.nodes() if n < len(G.nodes()) // 2]\n",
    "    group_2 = [n for n in G.nodes() if n >= len(G.nodes()) // 2]\n",
    "    \n",
    "    if not G0:\n",
    "        G0_edges = set()\n",
    "    else:\n",
    "        G0_edges = set(G0.edges())\n",
    "    G_edges = set(G.edges()) - G0_edges\n",
    "\n",
    "    G_group_1 = (set(nx.subgraph(G, group_1).edges()) & G_edges) - G0_edges\n",
    "    G0_group_1 = (set(nx.subgraph(G0, group_1).edges()))\n",
    "    G_group_2 = (set(nx.subgraph(G, group_2).edges()) & G_edges) - G0_edges\n",
    "    G0_group_2 = set(nx.subgraph(G0, group_2).edges())\n",
    "    G_between = G_edges - set(nx.subgraph(G, group_1).edges()) - set(nx.subgraph(G, group_2).edges()) - G0_edges\n",
    "    G0_between = G0_edges - set(nx.subgraph(G0, group_1).edges()) - set(nx.subgraph(G0, group_2).edges())\n",
    "    pos = nx.spring_layout(G)\n",
    "\n",
    "    if not use_netgraph:\n",
    "       \n",
    "        node_color = ['#c0392b' if n in group_1 else '#2980b9' for n in G.nodes()]\n",
    "\n",
    "        if not G0:\n",
    "            nx.draw(G, pos, ax=ax, node_size=10, width=0.5, node_color=node_color, alpha=0.7, edge_color='#34495e')\n",
    "        else:\n",
    "            \n",
    "            nx.draw_networkx_edges(G, pos, edgelist=G0_edges, width=0.5, alpha=0.5, edge_color='#34495e', ax=ax)\n",
    "            nx.draw_networkx_edges(G, pos, edgelist=G_between, width=1.0, alpha=1, edge_color='#f1c40f', ax=ax)\n",
    "            nx.draw_networkx_edges(G, pos, edgelist=G_group_1, width=2, alpha=1, edge_color='#c0392b', ax=ax)\n",
    "            nx.draw_networkx_edges(G, pos, edgelist=G0_group_1, width=1.0, alpha=0.5, edge_color='#e74c3c', ax=ax)\n",
    "\n",
    "            nx.draw_networkx_edges(G, pos, edgelist=G_group_2, width=2, alpha=1, edge_color='#2980b9', ax=ax)\n",
    "            nx.draw_networkx_edges(G, pos, edgelist=G0_group_2, width=1.0, alpha=0.5, edge_color='#3498db', ax=ax)\n",
    "\n",
    "            nx.draw_networkx_nodes(G, pos, nodelist=list(G.nodes()), node_size=10, node_color=node_color, alpha=0.7, ax=ax)\n",
    "    else:\n",
    "        node2community = {i: 0 if i < len(G.nodes()) // 2 else 1 for i in G.nodes()}\n",
    "        node_color = {i : '#c0392b' if node2community[i] == 0 else  '#2980b9' for i in G.nodes()}\n",
    "\n",
    "        edge_color = {}\n",
    "        edge_width = {}\n",
    "        edge_alpha = {}\n",
    "        for (u, v) in G.edges():\n",
    "            if (u, v) in G_group_1:\n",
    "                edge_color[u, v] = '#c0392b'\n",
    "            elif (u, v) in G_group_2:\n",
    "                edge_color[u, v] = '#2980b9'\n",
    "            elif (u, v) in G0_group_1:\n",
    "                edge_color[u, v] = '#e74c3c'\n",
    "            elif (u, v) in G0_group_2:\n",
    "                edge_color[u, v] = '#3498db'\n",
    "            elif (u, v) in G_between:\n",
    "                edge_color[u, v] = '#f1c40f'\n",
    "            else:\n",
    "                edge_color[u, v] = '#bdc3c7'\n",
    "\n",
    "            if (u, v) in G_group_1 or (u, v) in G_group_2 or (u, v) in G_between:\n",
    "                edge_width[u, v] = 2\n",
    "                edge_alpha[u, v] = 1\n",
    "            else:\n",
    "                edge_width[u, v] = 1\n",
    "                edge_alpha[u, v] = 0.5\n",
    "\n",
    "        # netgraph.Graph(G, node_layout='community', node_color=node_color, node_layout_kwargs=dict(node_to_community=node2community), node_size=2.5, edge_color=edge_color, edge_layout='bundled', edge_layout_kwargs=dict(k=2000), ax=ax)\n",
    "        netgraph.Graph(G, node_layout=pos, node_color=node_color, node_layout_kwargs=dict(node_to_community=node2community), node_size=2.5, edge_color=edge_color, edge_width=edge_width, edge_alpha=edge_alpha, ax=ax)\n",
    "\n",
    "\n",
    "    ax.set_axis_off()\n",
    "\n",
    "def prob_edge_wihin_community(G, G0):\n",
    "    group_1 = [n for n in G.nodes() if n < len(G.nodes()) // 2]\n",
    "    group_2 = [n for n in G.nodes() if n >= len(G.nodes()) // 2]\n",
    "    \n",
    "    G0_edges = set(G0.edges())\n",
    "\n",
    "    G_edges = set(G.edges()) - G0_edges\n",
    "\n",
    "    G_between = G_edges - set(nx.subgraph(G, group_1).edges()) - set(nx.subgraph(G, group_2).edges()) - G0_edges\n",
    "\n",
    "    try:\n",
    "        return 1 - len(G_between) / (1e-1 + len(G_edges))\n",
    "    except:\n",
    "        return 0\n",
    "\n",
    "def analyze_experiments(filename, num_common_neighbors=True):\n",
    "\n",
    "    suffix = os.path.split(os.path.splitext(filename)[0])[-1]\n",
    "\n",
    "    with open(filename) as f:\n",
    "        lines = f.read().splitlines()\n",
    "\n",
    "    data = []\n",
    "\n",
    "    for line in lines:\n",
    "        data.append(json.loads(line))\n",
    "\n",
    "    transitivities = collections.defaultdict(list)\n",
    "    algebraic_connectivities = collections.defaultdict(list)\n",
    "    probabilities_of_edge_within_community = collections.defaultdict(list)\n",
    "    # partition_qualitys = collections.defaultdict(list)\n",
    "\n",
    "    final_graphs = collections.defaultdict(list)\n",
    "\n",
    "    for d in data:\n",
    "        Gs = []\n",
    "        for i, graph in enumerate(d['graphs']):\n",
    "            G = nx.Graph()\n",
    "\n",
    "            for k, v in graph.items():\n",
    "                k = int(k)\n",
    "                G.add_node(k)\n",
    "                for n in v:\n",
    "                    G.add_edge(k, n)\n",
    "\n",
    "            G.remove_edges_from(nx.selfloop_edges(G))\n",
    "\n",
    "            if i > 0:\n",
    "                print('new edge', set(G.edges()) - set(Gs[0].edges()))\n",
    "\n",
    "            Gs.append(G)\n",
    "\n",
    "        fig, ax = plt.subplots(1, 4, figsize=(20, 5))\n",
    "\n",
    "        fig.suptitle(f'Temperature = {d[\"temperature\"]}')\n",
    "\n",
    "        for i, t in enumerate([0, len(Gs) // 2, len(Gs) - 1]):\n",
    "            G = Gs[t]\n",
    "            ax[i].set_title(f'$t = {t}$')\n",
    "            draw_graph(G, ax=ax[i], G0=Gs[0])\n",
    "\n",
    "            # print(d['reasons'])\n",
    "\n",
    "        group_1 = [n for n in G.nodes() if n < len(G.nodes()) // 2]\n",
    "        group_2 = [n for n in G.nodes() if n >= len(G.nodes()) // 2]\n",
    "\n",
    "\n",
    "        final_graphs[d['n'], d['temperature']].append((Gs[-1], Gs[0]))\n",
    "\n",
    "        initial_transitivity = nx.transitivity(Gs[0])\n",
    "\n",
    "        transitivity = [nx.transitivity(G) - initial_transitivity for G in Gs]\n",
    "\n",
    "        algebraic_connectivity = [nx.algebraic_connectivity(G) for G in Gs]\n",
    "\n",
    "        probability_of_edge_within_community = [prob_edge_wihin_community(G, Gs[0]) for G in Gs[1:]]\n",
    "\n",
    "        # partition_quality = [nx.community.partition_quality(G, communities)[0] for G in Gs]\n",
    "\n",
    "        # ax[-1].set_title('Metrics')\n",
    "        ax[-1].plot(transitivity, label='Marginal Transitivity', color='#c0392b')\n",
    "\n",
    "        ax_y = ax[-1].twinx()\n",
    "\n",
    "        ax_y.plot(algebraic_connectivity, label='Algebraic Connectivity', color='#2980b9')\n",
    "        ax[-1].set_xlabel('t')\n",
    "        ax[-1].set_ylabel('Transitivity', color='#c0392b')\n",
    "        ax_y.set_ylabel('Algebraic Connectivity', color='#2980b9')\n",
    "\n",
    "        transitivities[d['n'], d['temperature']].append(transitivity)\n",
    "        algebraic_connectivities[d['n'], d['temperature']].append(algebraic_connectivity)\n",
    "        probabilities_of_edge_within_community[d['n'], d['temperature']].append(probability_of_edge_within_community)\n",
    "        # partition_qualitys[d['n'], d['temperature']].append(partition_quality)\n",
    "\n",
    "        fig.tight_layout()\n",
    "        fig.savefig(f'figures/principle_2/{suffix}_{d[\"n\"]}_{d[\"simulation\"]}_{d[\"temperature\"]}{\"_neighbors\" if not num_common_neighbors else \"\"}.png', dpi=300)\n",
    "\n",
    "    palette = ['#e67e22', '#f1c40f', '#7f8c8d', '#c0392b', '#2980b9', '#34495e']\n",
    "\n",
    "\n",
    "    fig, ax = plt.subplots(4, len(transitivities), figsize=(5 * len(transitivities), 10), squeeze=False, sharey='row')\n",
    "\n",
    "    fig_combined, ax_combined = plt.subplots(1, 3, figsize=(10, 5), squeeze=False)\n",
    "\n",
    "    ax_combined[0, 0].set_ylabel('Transitivity')\n",
    "    ax_combined[0, 1].set_ylabel('Algebraic Connectivity')\n",
    "    ax_combined[0, 2].set_ylabel('Probability of Edge within Community')\n",
    "    # ax_combined[0, 3].set_ylabel('Partition Quality')\n",
    "    ax_combined[0, 0].set_xlabel('t')\n",
    "    ax_combined[0, 1].set_xlabel('t')\n",
    "    ax_combined[0, 2].set_xlabel('t')\n",
    "    # ax_combined[0, 3].set_xlabel('t')\n",
    "\n",
    "    ax_combined[0, 0].set_xlim(0, 49)\n",
    "    ax_combined[0, 1].set_xlim(0, 49)\n",
    "    ax_combined[0, 2].set_xlim(0, 49)\n",
    "    # ax_combined[0, 3].set_xlim(0, 49)\n",
    "\n",
    "    fig_final, ax_final = plt.subplots(1, len(final_graphs) + 2, figsize=(5 * (2 + len(final_graphs)), 5), squeeze=False, gridspec_kw={'width_ratios': [1] * len(final_graphs) + [0.5, 0.5]})\n",
    "\n",
    "    ax_final[0, -1].spines[['right', 'top']].set_visible(False)\n",
    "    ax_final[0, -2].spines[['right', 'top']].set_visible(False)\n",
    "\n",
    "\n",
    "\n",
    "    for i, (k, v) in enumerate(sorted(final_graphs.items())):\n",
    "        G, G0 = v[0]\n",
    "        draw_graph(G, ax=ax_final[0, i], G0=G0)\n",
    "\n",
    "        ax_final[0, i].set_title(f'Temperature = {k[1]}')\n",
    "\n",
    "\n",
    "    # ax_final[0, -1].set_ylabel('Algebraic Connectivity')\n",
    "    # ax_final[0, -1].set_xticks([])\n",
    "\n",
    "    ax_final[0, -2].set_ylabel('Marginal Transitivity')\n",
    "    ax_final[0, -2].set_xticks([])\n",
    "        \n",
    "    ax_final[0, -1].set_ylabel('Pr. of Edge within Community')\n",
    "    ax_final[0, -1].set_xticks([])\n",
    "\n",
    "    for i, (k, c) in enumerate(zip(sorted(transitivities.keys()), palette)):\n",
    "        v = transitivities[k]\n",
    "        v = np.array(v)\n",
    "\n",
    "        mean = v.mean(axis=0)\n",
    "        std = v.std(axis=0)\n",
    "\n",
    "        ci = 1.96 * std / np.sqrt(len(v))\n",
    "\n",
    "        ax[0, i].plot(mean, color='#2980b9', label='LLM')\n",
    "        ax[0, i].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#2980b9')\n",
    "\n",
    "        ax[0, i].set_title(f'Temperature = {k[1]}')\n",
    "\n",
    "        ax[0, i].set_xlabel('t')\n",
    "\n",
    "        if i == 0:\n",
    "            ax[0, i].set_ylabel('Marginal Transitivity')\n",
    "\n",
    "        ax[0, i].set_ylim(0.35, 0.5)\n",
    "\n",
    "        ax[0, i].set_xlim(0, len(mean) - 1)\n",
    "\n",
    "        ax_combined[0, 0].plot(mean, color=c, label='Temp = ' + str(k[1]))\n",
    "        ax_combined[0, 0].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color=c)\n",
    "\n",
    "        ax_final[0, -2].bar(i, mean[-1], color=palette[i], alpha=0.6, label='Temp = ' + str(k[1]))\n",
    "        ax_final[0, -2].errorbar(i, mean[-1], yerr=ci[-1], color='black', alpha=1)\n",
    "\n",
    "    ax_combined[0, 0].set_ylim(0.35, 0.5)\n",
    "\n",
    "\n",
    "    for i, (k, c) in enumerate(zip(sorted(algebraic_connectivities.keys()), palette)):\n",
    "        v = algebraic_connectivities[k]\n",
    "        v = np.array(v)\n",
    "\n",
    "        mean = v.mean(axis=0)\n",
    "        std = v.std(axis=0)\n",
    "\n",
    "        ci = 1.96 * std / np.sqrt(len(v))\n",
    "\n",
    "\n",
    "        ax[1, i].plot(mean, color='#2980b9', label='LLM')\n",
    "        ax[1, i].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#2980b9')\n",
    "        ax[1, i].set_ylim(3.5, 6)\n",
    "\n",
    "        ax[1, i].set_xlim(0, len(mean) - 1)\n",
    "\n",
    "        if i == 0:\n",
    "            ax[1, i].set_ylabel('Algebraic Connectivity')\n",
    "\n",
    "        ax_combined[0, 1].plot(mean, color=c, label='Temp = ' + str(k[1]))\n",
    "        ax_combined[0, 1].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color=c)\n",
    "\n",
    "        # ax_final[0, -1].bar(i, mean[-1], color=palette[i], alpha=0.5, label='Temp = ' + str(k[1]))\n",
    "        # ax_final[0, -1].errorbar(i, mean[-1], yerr=ci[-1], color=palette[i], alpha=0.5)\n",
    "\n",
    "    for i, (k, c) in enumerate(zip(sorted(probabilities_of_edge_within_community.keys()), palette)):\n",
    "        v = probabilities_of_edge_within_community[k]\n",
    "        v = np.array(v)\n",
    "\n",
    "        mean = v.mean(axis=0)\n",
    "        std = v.std(axis=0)\n",
    "\n",
    "        ci = 1.96 * std / np.sqrt(len(v))\n",
    "\n",
    "        # ax[2, i].plot(mean, color='#2980b9', label='LLM')\n",
    "        # ax[2, i].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#2980b9')\n",
    "\n",
    "        # ax[2, i].set_title(f'Temperature = {k[1]}')\n",
    "\n",
    "        # ax[2, i].set_xlabel('t')\n",
    "\n",
    "        # if i == 0:\n",
    "        #     ax[2, i].set_ylabel('Probability of Edge within Community')\n",
    "\n",
    "        # ax[2, i].set_ylim(0.35, 0.5)\n",
    "\n",
    "        # ax[2, i].set_xlim(0, len(mean) - 1)\n",
    "\n",
    "        # ax_combined[0, 2].plot(mean, color=c, label='Temp = ' + str(k[1]))\n",
    "        # ax_combined[0, 2].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color=c)\n",
    "\n",
    "        ax_final[0, -1].bar(i, mean[-1], color=palette[i], alpha=0.5, label='Temp = ' + str(k[1]))\n",
    "        ax_final[0, -1].errorbar(i, mean[-1], yerr=ci[-1], color='black', alpha=0.5)\n",
    "\n",
    "\n",
    "    ax_combined[0, 2].set_ylim(0, 1)\n",
    "\n",
    "    # Null models\n",
    "    transitivities_null = { 'random' : collections.defaultdict(list), 'winner' : collections.defaultdict(list), 'sbm' : collections.defaultdict(list) }\n",
    "    algebraic_connectivities_null = { 'random' : collections.defaultdict(list), 'winner' : collections.defaultdict(list), 'sbm' : collections.defaultdict(list) }\n",
    "    probabilities_of_edge_within_community_null = { 'random' : collections.defaultdict(list), 'winner' : collections.defaultdict(list), 'sbm' : collections.defaultdict(list) }\n",
    "    # partition_qualitys_null = { 'random' : collections.defaultdict(list), 'winner' : collections.defaultdict(list) }\n",
    "\n",
    "    for d in data:\n",
    "        for method in ['random']:\n",
    "            Gs, _ = network_growth(d['n'], d['temperature'], method=method)\n",
    "\n",
    "            initial_transitivity = nx.transitivity(Gs[0])\n",
    "\n",
    "            transitivity = [nx.transitivity(G) - initial_transitivity for G in Gs]\n",
    "\n",
    "            transitivities_null[method][d['n'], d['temperature']].append(transitivity)\n",
    "\n",
    "\n",
    "            algebraic_connectivity = [nx.algebraic_connectivity(G) for G in Gs]\n",
    "\n",
    "            algebraic_connectivities_null[method][d['n'], d['temperature']].append(algebraic_connectivity)\n",
    "\n",
    "            group_1 = [n for n in G.nodes() if n < len(G.nodes()) // 2]\n",
    "            group_2 = [n for n in G.nodes() if n >= len(G.nodes()) // 2]\n",
    "\n",
    "            communities = [group_1, group_2]\n",
    "\n",
    "            probability_of_edge_within_community = [prob_edge_wihin_community(G, Gs[0]) for G in Gs[1:]]\n",
    "\n",
    "            probabilities_of_edge_within_community_null[method][d['n'], d['temperature']].append(probability_of_edge_within_community)\n",
    "\n",
    "            # partition_quality = [nx.algorithms.community.partition_quality(G, communities)[0] for G in Gs]\n",
    "            # partition_qualitys_null[method][d['n'], d['temperature']].append(partition_quality)\n",
    "    \n",
    "\n",
    "    for j, method in enumerate(['random']):\n",
    "        for i, (k, v) in enumerate(transitivities_null[method].items()):\n",
    "            v = np.array(v)\n",
    "\n",
    "            mean = v.mean(axis=0)\n",
    "            std = v.std(axis=0)\n",
    "\n",
    "            ci = 1.96 * std / np.sqrt(len(v))\n",
    "\n",
    "            ax[0, i].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':', label=method.capitalize())\n",
    "            ax[0, i].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e')\n",
    "\n",
    "            if i == 0:\n",
    "                ax_combined[0, 0].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':', label=method.capitalize())\n",
    "                ax_combined[0, 0].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e')\n",
    "\n",
    "            if i == 0:\n",
    "                ax_final[0, -2].bar(j + 3, mean[-1], color=palette[j+3], alpha=0.6, label=method.capitalize())\n",
    "                ax_final[0, -2].errorbar(j + 3, mean[-1], yerr=ci[-1], color='black', alpha=1)\n",
    "            \n",
    "            print('Transitivity T-test', k, method, scipy.stats.ttest_ind([x[-1] for x in transitivities[k]], [x[-1] for x in transitivities_null[method][k]], equal_var=False))\n",
    "\n",
    "        for i, (k, v) in enumerate(algebraic_connectivities_null[method].items()):\n",
    "            v = np.array(v)\n",
    "\n",
    "            mean = v.mean(axis=0)\n",
    "            std = v.std(axis=0)\n",
    "\n",
    "            ci = 1.96 * std / np.sqrt(len(v))\n",
    "\n",
    "            ax[1, i].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':', label=method.capitalize())\n",
    "            ax[1, i].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e')\n",
    "            if i == 0:\n",
    "                ax_combined[0, 1].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':')\n",
    "                ax_combined[0, 1].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e', hatch='||')\n",
    "\n",
    "            # if i == 0:\n",
    "            #     ax_final[0, -1].bar(j + 3, mean[-1], color=palette[j+3], alpha=0.2, label=method.capitalize())\n",
    "            #     ax_final[0, -1].errorbar(j + 3, mean[-1], yerr=ci[-1], color=palette[j+3], alpha=0.5)\n",
    "\n",
    "            print('Algebraic Connectivity T-test', k, method, scipy.stats.ttest_ind([x[-1] for x in algebraic_connectivities[k]], [x[-1] for x in algebraic_connectivities_null[method][k]], equal_var=False))\n",
    "\n",
    "        for i, (k, v) in enumerate(probabilities_of_edge_within_community_null[method].items()):\n",
    "            \n",
    "\n",
    "            v = np.array(v)\n",
    "\n",
    "            mean = v.mean(axis=0)\n",
    "            std = v.std(axis=0)\n",
    "\n",
    "            ci = 1.96 * std / np.sqrt(len(v))\n",
    "\n",
    "            ax[1, i].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':', label=method.capitalize())\n",
    "            ax[1, i].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e')\n",
    "            if i == 0:\n",
    "                ax_combined[0, 2].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':')\n",
    "                ax_combined[0, 2].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e', hatch='||')\n",
    "\n",
    "            if i == 0:                \n",
    "                ax_final[0, -1].bar(j + 3, mean[-1], color=palette[j+3], alpha=0.6, label=method.capitalize())\n",
    "                ax_final[0, -1].errorbar(j + 3, mean[-1], yerr=ci[-1], color='black', alpha=1)\n",
    "\n",
    "            print('Probability of edge within community T-test', k, method, scipy.stats.ttest_ind([x[-1] for x in probabilities_of_edge_within_community[k]], [x[-1] for x in probabilities_of_edge_within_community_null[method][k]], equal_var=False))\n",
    "\n",
    "        # for i, (k, v) in enumerate(probabilities_of_edge_within_community_null[method].items()):\n",
    "        #     v = np.array(v)\n",
    "\n",
    "        #     mean = v.mean(axis=0)\n",
    "        #     std = v.std(axis=0)\n",
    "\n",
    "        #     ci = 1.96 * std / np.sqrt(len(v))\n",
    "\n",
    "        #     ax[1, i].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':', label=method.capitalize())\n",
    "        #     ax[1, i].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e')\n",
    "        #     if i == 0:\n",
    "        #         ax_combined[0, 3].plot(mean, color='#c0392b' if method == 'random' else '#34495e', linestyle='--' if method == 'random' else ':')\n",
    "        #         ax_combined[0, 3].fill_between(np.arange(len(mean)), mean - ci, mean + ci, alpha=0.2, color='#c0392b' if method == 'random' else '#34495e', hatch='||')\n",
    "\n",
    "        #     if i == 0:\n",
    "        #         ax_final[0, -2].bar(j + 3, mean[-1], color=palette[j+3], alpha=0.6, label=method.capitalize())\n",
    "        #         ax_final[0, -2].errorbar(j + 3, mean[-1], yerr=ci[-1], color='black', alpha=1)\n",
    "\n",
    "        #     print('Partition Quality T-test', k, method, scipy.stats.ttest_ind([x[-1] for x in partition_qualitys[k]], [x[-1] for x in partition_qualitys_null[method][k]], equal_var=False))\n",
    "\n",
    "    ax_final[0, -1].legend(bbox_to_anchor=(1, 0.5), loc='center left', frameon=False)\n",
    "\n",
    "    ax[0, 0].legend(loc='upper left')\n",
    "    ax[1, 0].legend(loc='upper left')\n",
    "\n",
    "    ax_combined[0, 0].legend(loc='upper left')\n",
    "    # ax_combined[0, 1].legend(loc='upper left')\n",
    "\n",
    "    fig_combined.tight_layout()\n",
    "\n",
    "    fig.tight_layout()\n",
    "\n",
    "    fig.savefig(f'figures/principle_2/{suffix}_overall{\"_neighbors\" if not num_common_neighbors else \"\"}.png', dpi=300)\n",
    "\n",
    "    fig_combined.savefig(f'figures/principle_2/{suffix}_overall_combined{\"_neighbors\" if not num_common_neighbors else \"\"}.png', dpi=300)\n",
    "     \n",
    "\n",
    "    fig_final.tight_layout()\n",
    "\n",
    "    fig_final.savefig(f'figures/principle_2/{suffix}_final{\"_neighbors\" if not num_common_neighbors else \"\"}.png', dpi=300)\n",
    "\n",
    "\n",
    "def get_table(filenames, sfx='', environments=True):\n",
    "\n",
    "    records = []\n",
    "\n",
    "    num_graphs = 0\n",
    "\n",
    "    for filename in filenames:\n",
    "        print(filename)\n",
    "        suffix = os.path.split(os.path.splitext(filename)[0])[-1]\n",
    "        suffix = suffix.split('+')\n",
    "\n",
    "        if len(suffix) == 3:\n",
    "            model = suffix[-2]\n",
    "            environment = suffix[-1]\n",
    "        elif len(suffix) == 2:\n",
    "            model = suffix[-1]\n",
    "            environment = 'Baseline'\n",
    "        else:\n",
    "            environment = 'Baseline'\n",
    "\n",
    "        with open(filename) as f:\n",
    "            lines = f.read().splitlines()\n",
    "\n",
    "        data = []\n",
    "\n",
    "        for line in lines:\n",
    "            data.append(json.loads(line))\n",
    "\n",
    "        for d in data:\n",
    "            Gs = []\n",
    "\n",
    "            for i, graph in enumerate(d['graphs']):\n",
    "                G = nx.Graph()\n",
    "\n",
    "                for k, v in graph.items():\n",
    "                    k = int(k)\n",
    "                    G.add_node(k)\n",
    "                    for n in v:\n",
    "                        G.add_edge(k, n)\n",
    "\n",
    "                if i == 0:\n",
    "                    top_common_neighbors = np.zeros(len(G.nodes()))\n",
    "                    total = 0\n",
    "                else:\n",
    "                    new_edge =  set(G.edges()) - set(Gs[-1].edges())\n",
    "                    if len(new_edge) == 0:\n",
    "                        continue\n",
    "\n",
    "                    new_edge = new_edge.pop()\n",
    "\n",
    "                    u, v = new_edge\n",
    "\n",
    "                    common_neighbors_u = [len(set(G.neighbors(u)) & set(G.neighbors(n))) for n in G.nodes() if n != u]\n",
    "\n",
    "                    # find what position the new edge is in the sorted list of common neighbors\n",
    "                    try:\n",
    "                        pos = np.argsort(common_neighbors_u)[::-1].tolist().index(v)\n",
    "                        top_common_neighbors[pos] += 1\n",
    "                        total += 1\n",
    "                    except:\n",
    "                        pass\n",
    "\n",
    "                G.remove_edges_from(nx.selfloop_edges(G))\n",
    "\n",
    "                Gs.append(G)\n",
    "\n",
    "            top_common_neighbors /= total\n",
    "\n",
    "            top_common_neighbors = np.cumsum(top_common_neighbors)\n",
    "            top_common_neighbors = np.insert(top_common_neighbors, 0, 0)\n",
    "\n",
    "\n",
    "            initial_transitivity = nx.transitivity(Gs[0])\n",
    "            final_transitivity = nx.transitivity(Gs[-1])\n",
    "            marginal_transitivity = final_transitivity - initial_transitivity\n",
    "            final_probability_of_edge_within_community = prob_edge_wihin_community(Gs[-1], Gs[1])\n",
    "\n",
    "            record = {\n",
    "                'Model' : model,\n",
    "                'Environment' : environment,\n",
    "                'Temperature' : d['temperature'],\n",
    "                'Marginal Transitivity' : marginal_transitivity,\n",
    "                # 'Algebraic Connectivity' : final_algebraic_connectivity,\n",
    "                'Prob. of Edge within Community' : final_probability_of_edge_within_community,\n",
    "                'Probability of Connecting to Top-$k$' : top_common_neighbors,\n",
    "                'Top-$k$' : np.arange(0, len(top_common_neighbors)) / len(top_common_neighbors)\n",
    "            }\n",
    "\n",
    "            records.append(record)\n",
    "\n",
    "        \n",
    "\n",
    "    df = pd.DataFrame(records)\n",
    "\n",
    "    rename_models = {\n",
    "        'gpt-3.5-turbo' : 'GPT-3.5',\n",
    "        'gpt-4o-mini' : 'GPT-4 Mini',\n",
    "        'meta-meta-llama-3-70b-instruct' : 'LLAMA-3',\n",
    "        'claude-3-5-sonnet-20240620' : 'Claude 3.5',\n",
    "        'gpt-3.5-turbo_cot' : 'GPT-3.5',\n",
    "        'gpt-4o-mini_cot' : 'GPT-4 Mini',\n",
    "        'meta-meta-llama-3-70b-instruct_cot' : 'LLAMA-3',\n",
    "        'claude-3-5-sonnet-20240620_cot' : 'Claude 3.5'\n",
    "    }\n",
    "\n",
    "    rename_env = {\n",
    "        'school' : 'School',\n",
    "        'work' : 'Work',\n",
    "        'community' : 'Community',\n",
    "        'school_cot' : 'School',\n",
    "        'work_cot' : 'Work',\n",
    "        'community_cot' : 'Community',\n",
    "    }\n",
    "\n",
    "\n",
    "    ncols = 2 + int(environments)\n",
    "\n",
    "    df['Model'] = df['Model'].apply(lambda x: rename_models.get(x, x))\n",
    "    df['Environment'] = df['Environment'].apply(lambda x: rename_env.get(x, x))\n",
    "\n",
    "    df_model = df.query('Environment == \"Baseline\" and Temperature == 0.5')\n",
    "    df_environment = df.query('Model == \"GPT-4 Mini\"')\n",
    "\n",
    "    df_temperature = df.query('Model == \"GPT-4 Mini\" and Environment == \"Baseline\"')\n",
    "\n",
    "    fig, ax = plt.subplots(2, ncols, figsize=(5 * ncols, 10))\n",
    "    \n",
    "    sc_model = sns.barplot(data=df_model, y='Marginal Transitivity', x='Model', ax=ax[0, 0], palette=['#e67e22', '#f1c40f', '#3498db', '#7f8c8d', '#c0392b', '#34495e', '#2980b9'])\n",
    "\n",
    "    sc_temperature = sns.barplot(data=df_temperature, y='Marginal Transitivity', x='Temperature', ax=ax[0, 1], palette=['#e67e22', '#f1c40f', '#3498db', '#7f8c8d', '#c0392b', '#34495e', '#2980b9'])\n",
    "    \n",
    "    ax[0, 0].set_ylabel('D')\n",
    "\n",
    "    ax[0, 0].set_xticks([])\n",
    "    ax[0, 1].set_xticks([])\n",
    "\n",
    "    ax[0, 0].set_ylim(0, 0.1)\n",
    "    ax[0, 1].set_ylim(0, 0.1)\n",
    "\n",
    "    ax[0, 0].set_xlabel('')\n",
    "    ax[0, 1].set_xlabel('')\n",
    "\n",
    "\n",
    "    ax[0, 1].get_yaxis().set_visible(False)\n",
    "\n",
    "    if environments:\n",
    "        sc_environment = sns.barplot(data=df_environment, y='Marginal Transitivity', x='Environment', ax=ax[0, ncols-1], palette=['#e67e22', '#f1c40f', '#3498db', '#7f8c8d', '#c0392b', '#34495e', '#2980b9'])\n",
    "        ax[0, 2].get_yaxis().set_visible(False)\n",
    "        ax[0, 2].set_xlabel('')\n",
    "        ax[0, 2].set_ylim(0, 0.1)\n",
    "        ax[0, 2].set_xticks([])\n",
    "\n",
    "    sc_model = sns.barplot(data=df_model, y='Prob. of Edge within Community', x='Model', ax=ax[1, 0], palette=['#e67e22', '#f1c40f', '#3498db', '#7f8c8d', '#c0392b', '#34495e', '#2980b9'])\n",
    "    sc_temperature = sns.barplot(data=df_temperature, y='Prob. of Edge within Community', x='Temperature', ax=ax[1, 1], palette=['#e67e22', '#f1c40f', '#3498db', '#7f8c8d', '#c0392b', '#34495e', '#2980b9'])\n",
    "\n",
    "    sc_model.set_xticklabels(sc_model.get_xticklabels(), rotation=90)\n",
    "    sc_temperature.set_xticklabels(sc_temperature.get_xticklabels(), rotation=90)\n",
    "\n",
    "    ax[1, 0].set_ylabel('$\\\\hat p$')\n",
    "    ax[1, 1].set_ylabel('')\n",
    "\n",
    "    ax[1, 0].set_xlabel('')\n",
    "    ax[1, 1].set_xlabel('')\n",
    "\n",
    "    ax[0, 0].set_title('Model')\n",
    "    ax[0, 1].set_title('Temperature')\n",
    "\n",
    "    ax[1, 1].get_yaxis().set_visible(False)\n",
    "\n",
    "    ax[1, 0].set_ylim(0, 1)\n",
    "    ax[1, 1].set_ylim(0, 1)\n",
    "    \n",
    "    if environments:\n",
    "        sc_environment = sns.barplot(data=df_environment, y='Prob. of Edge within Community', x='Environment', ax=ax[1, 2], palette=['#e67e22', '#f1c40f', '#3498db', '#7f8c8d', '#c0392b', '#34495e', '#2980b9'])\n",
    "        sc_environment.set_xticklabels(sc_environment.get_xticklabels(), rotation=90)\n",
    "        ax[1, 2].set_ylim(0, 1)\n",
    "        ax[1, 2].get_yaxis().set_visible(False)\n",
    "        ax[1, 2].set_xlabel('')\n",
    "        ax[1, 2].set_ylabel('')\n",
    "        ax[1, 2].spines[['right', 'top']].set_visible(False)\n",
    "        ax[0, 2].spines[['right', 'top']].set_visible(False)\n",
    "        ax[0, 2].set_title('Environment')\n",
    "\n",
    "\n",
    "    ax[0, 0].spines[['right', 'top']].set_visible(False)\n",
    "    ax[0, 1].spines[['right', 'top']].set_visible(False)\n",
    "    ax[1, 0].spines[['right', 'top']].set_visible(False)\n",
    "    ax[1, 1].spines[['right', 'top']].set_visible(False)\n",
    "\n",
    "    fig.savefig(f'figures/triadic_closure{sfx}.png', dpi=300, bbox_inches='tight')\n",
    "\n",
    "    fig, ax = plt.subplots(1, ncols, figsize=(5 * ncols, 5))\n",
    "\n",
    "    palette=['#e67e22', '#f1c40f', '#3498db', '#7f8c8d', '#c0392b', '#34495e', '#2980b9']\n",
    "\n",
    "    fig.suptitle('Probability of Connecting to Top-$k$ Common Neighbors', fontsize=SMALL_SIZE)\n",
    "\n",
    "    breakpoints = np.array([0.1, 0.2, 0.3, 0.4, 0.5])\n",
    "\n",
    "    breakpoint_max = np.max(breakpoints)\n",
    "    breakpoint_min = np.min(breakpoints)\n",
    "\n",
    "    for i in range(len(ax)):\n",
    "        ax[i].plot([0, 100 * breakpoint_max], [0, breakpoint_max], color='black', linestyle='--')\n",
    "        ax[i].set_xlim(100 * breakpoint_min, 100 * breakpoint_max)\n",
    "\n",
    "    for i, model in enumerate(df_model['Model'].unique()):\n",
    "        temp = df_model[df_model['Model'] == model]\n",
    "\n",
    "        n = len(temp['Top-$k$'].values[0])\n",
    "\n",
    "        indices = np.array([int(x * n) for x in breakpoints])\n",
    "\n",
    "        color = palette[i]\n",
    "        linewidth = 1\n",
    "\n",
    "        if model == 'GPT-4 Mini' and df_model['Environment'].values[i] == 'Baseline' and df_model['Temperature'].values[i] == 0.5:\n",
    "            color = '#34495e'\n",
    "            linewidth = 3\n",
    "\n",
    "        ax[0].plot(100 * temp['Top-$k$'].values[0][indices], temp['Probability of Connecting to Top-$k$'].values.mean(0)[indices], label=model, color=color, linewidth=linewidth, marker='x')\n",
    "\n",
    "    ax[0].set_title('Model')\n",
    "    ax[0].set_xlabel('Top-$k$ (%)')\n",
    "    # ax[0].set_xscale('log')\n",
    "    # ax[0].set_yscale('log')\n",
    "    ax[0].set_ylabel('')\n",
    "    \n",
    "    ax[0].legend(fontsize=0.7*SMALL_SIZE, loc='upper left')\n",
    "\n",
    "    for i, temperature in enumerate(df_temperature['Temperature'].unique()):\n",
    "        temp = df_temperature[df_temperature['Temperature'] == temperature]\n",
    "\n",
    "        n = len(temp['Top-$k$'].values[0])\n",
    "\n",
    "        indices = np.array([int(x * n) for x in breakpoints])\n",
    "\n",
    "        color = palette[i]\n",
    "        linewidth = 1\n",
    "\n",
    "        if temperature == 0.5 and df_temperature['Model'].values[i] == 'GPT-4 Mini' and df_temperature['Environment'].values[i] == 'Baseline':\n",
    "            color = '#34495e'\n",
    "            linewidth = 3\n",
    "            \n",
    "\n",
    "        ax[1].plot(100 * temp['Top-$k$'].values[0][indices], temp['Probability of Connecting to Top-$k$'].values.mean(0)[indices], label=f'{temperature}', color=color, linewidth=linewidth, marker='x')\n",
    "\n",
    "\n",
    "    ax[1].set_title('Temperature')\n",
    "    ax[1].set_xlabel('Top-$k$ (%)')\n",
    "    # ax[1].set_xscale('log')\n",
    "    # ax[1].set_yscale('log')\n",
    "    ax[1].set_ylabel('')\n",
    "\n",
    "    ax[1].legend(fontsize=0.7*SMALL_SIZE, loc='upper left')\n",
    "\n",
    "    if environments:\n",
    "\n",
    "        for i, environment in enumerate(df_environment['Environment'].unique()):\n",
    "            temp = df_environment[df_environment['Environment'] == environment]\n",
    "\n",
    "            n = len(temp['Top-$k$'].values[0])\n",
    "\n",
    "            indices = np.array([int(x * n) for x in breakpoints])\n",
    "\n",
    "            color = palette[i]\n",
    "            linewidth = 1\n",
    "\n",
    "            if environment == 'Baseline' and df_environment['Model'].values[i] == 'GPT-4 Mini' and df_environment['Temperature'].values[i] == 0.5:\n",
    "                color = '#34495e'\n",
    "                linewidth = 3        \n",
    "                ax[2].plot(100 * temp['Top-$k$'].values[0][indices], df_model[df_model['Model'] == 'GPT-4 Mini']['Probability of Connecting to Top-$k$'].values.mean(0)[indices], label='GPT-4 Mini', color=color, linewidth=linewidth, marker='x')\n",
    "\n",
    "            else:\n",
    "                ax[2].plot(100 * temp['Top-$k$'].values[0][indices], temp['Probability of Connecting to Top-$k$'].values.mean(0)[indices], label=environment, color=color, linewidth=linewidth, marker='x')\n",
    "\n",
    "\n",
    "        ax[2].set_title('Environment')\n",
    "        ax[2].set_xlabel('Top-$k$ (%)')\n",
    "    \n",
    "        ax[2].set_ylabel('')\n",
    "        ax[2].legend(fontsize=0.7*SMALL_SIZE, loc='upper left')\n",
    "        ax[2].set_ylim(0, 1)\n",
    "        ax[2].get_yaxis().set_visible(False)\n",
    "        ax[2].spines[['right', 'top']].set_visible(False)\n",
    "\n",
    "\n",
    "    ax[0].set_ylim(0, 1)\n",
    "    ax[1].set_ylim(0, 1)\n",
    "\n",
    "    # hide y axis numbers\n",
    "    ax[1].get_yaxis().set_visible(False)\n",
    "\n",
    "    ax[0].spines[['right', 'top']].set_visible(False)\n",
    "    ax[1].spines[['right', 'top']].set_visible(False)\n",
    "\n",
    "    fig.tight_layout()\n",
    "\n",
    "                     \n",
    "    fig.savefig(f'figures/top_kcommon{sfx}.png', dpi=300, bbox_inches='tight')\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_outfiles = []\n",
    "# Other environments\n",
    "for environment, role in zip(['school', 'work', 'community'], ['classmates', 'colleagues', 'neighbors']):\n",
    "    outfile = f'outputs/principle_2_neighbors+gpt-4o-mini+{environment}.jsonl'\n",
    "    run_network_formation_experiment(50, 50, 1, 10, outfile, [0.5], environment=environment, role=role, method='llm', num_common_neighbors=False, model='gpt-4o-mini')\n",
    "\n",
    "    table_outfiles.append(outfile)\n",
    "\n",
    "for model in ['gpt-3.5-turbo', 'meta/meta-llama-3-70b-instruct', 'gpt-4o-mini', 'claude-3-5-sonnet-20240620']:\n",
    "    outfile = f'outputs/principle_2_neighbors+{model.replace(\"/\", \"-\")}.jsonl'\n",
    "    if model == 'claude-3-5-sonnet-20240620':\n",
    "        run_network_formation_experiment(50, 50, 1, 10, outfile, [0.5, 1.0], method='llm', num_common_neighbors=False, model=model)\n",
    "    else:\n",
    "        run_network_formation_experiment(50, 50, 1, 10, outfile, [0.5, 1.0, 1.5], method='llm', num_common_neighbors=False, model=model)\n",
    "\n",
    "    table_outfiles.append(outfile)\n",
    "\n",
    "get_table(table_outfiles)\n",
    "\n",
    "# run_network_formation_experiment(50, 50, 1, 10, 'outputs/principle_2_neighbors.jsonl', [0.5, 1.0, 1.5], method='llm', num_common_neighbors=False)\n",
    "# analyze_experiments('outputs/principle_2_neighbors.jsonl', num_common_neighbors=False)\n",
    "\n",
    "# run_network_formation_experiment(50, 50, 1, 10, 'outputs/principle_2.jsonl', [0.5, 1.0, 1.5], method='llm',  num_common_neighbors=True)\n",
    "# analyze_experiments('outputs/principle_2.jsonl', num_common_neighbors=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "table_outfiles_cot = []\n",
    "\n",
    "# Other environments\n",
    "for environment, role in zip(['school', 'work', 'community'], ['classmates', 'colleagues', 'neighbors']):\n",
    "    outfile = f'outputs/principle_2_neighbors+gpt-4o-mini+{environment}_cot.jsonl'\n",
    "    run_network_formation_experiment(50, 50, 1, 5, outfile, [0.5, 1.0, 1.5], environment=environment, role=role, method='llm', num_common_neighbors=False, model='gpt-4o-mini', cot=True)\n",
    "\n",
    "    table_outfiles_cot.append(outfile)\n",
    "\n",
    "\n",
    "for model in ['gpt-3.5-turbo', 'meta/meta-llama-3-70b-instruct', 'gpt-4o-mini', 'claude-3-5-sonnet-20240620']:\n",
    "    outfile = f'outputs/principle_2_neighbors+{model.replace(\"/\", \"-\")}_cot.jsonl'\n",
    "\n",
    "    print(model)\n",
    "\n",
    "    if model == 'claude-3-5-sonnet-20240620':\n",
    "        run_network_formation_experiment(50, 50, 1, 5, outfile, [0.5, 1.0], method='llm', num_common_neighbors=False, model=model, cot=True)\n",
    "    else:\n",
    "        run_network_formation_experiment(50, 50, 1, 5, outfile, [0.5, 1.0, 1.5], method='llm', num_common_neighbors=False, model=model, cot=True)\n",
    "\n",
    "\n",
    "    table_outfiles_cot.append(outfile)\n",
    "\n",
    "get_table(table_outfiles_cot, sfx='_cot')\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}