cradle.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "def _2div(arr):\n",
    "    arr_res = arr.copy()\n",
    "    arr_pieces = []\n",
    "    pa = 0\n",
    "    st = 0\n",
    "    needdivcnt = 0\n",
    "    for i, a in enumerate(arr):\n",
    "        if a!=pa:\n",
    "            arr_pieces.append([st, i])\n",
    "            if (i-st)!=1: needdivcnt+=1\n",
    "            pa = a\n",
    "            st = i\n",
    "\n",
    "    arr_pieces.append([st, len(arr)])\n",
    "    if (len(arr)-st)!=1: needdivcnt+=1\n",
    "\n",
    "    offset = range(len(arr_pieces), len(arr_pieces)+needdivcnt)\n",
    "    p=0\n",
    "    for arr_p in arr_pieces:\n",
    "        length = arr_p[1] - arr_p[0]\n",
    "        if length == 1: continue\n",
    "        half_len = int(np.ceil(length / 2))\n",
    "        for j in range(arr_p[0]+half_len, arr_p[1]):\n",
    "            try:\n",
    "                arr_res[j] = offset[p]\n",
    "            except:\n",
    "                print('wtf')\n",
    "        p+=1\n",
    "    return arr_res\n",
    "\n",
    "def get_division_tree(n_agents):\n",
    "    agent2divitreeindex = np.arange(n_agents)\n",
    "    np.random.shuffle(agent2divitreeindex)\n",
    "    max_div = np.ceil(np.log2(n_agents)).astype(int)\n",
    "    levels = np.zeros(shape=(max_div+1, n_agents), dtype=int)\n",
    "    tree_of_agent = []*(max_div+1)\n",
    "    for ith, level in enumerate(levels):\n",
    "        if ith == 0: continue\n",
    "        res = _2div(levels[ith-1,:])\n",
    "        levels[ith,:] = res\n",
    "    res_levels = levels.copy()\n",
    "    for i, div_tree_index in enumerate(agent2divitreeindex):\n",
    "        res_levels[:, i] = levels[:, div_tree_index]\n",
    "    return res_levels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{0: 2475, 1: 2484, 2: 5041}, {0: 1938, 1: 5041, 2: 3021}, {0: 7031, 1: 1946, 2: 1023}]\n"
     ]
    }
   ],
   "source": [
    "from torch.distributions.categorical import Categorical\n",
    "probs = torch.Tensor([\n",
    "    [0.3, 0.3, 0.4],\n",
    "    [0.2, 0.5, 0.3],\n",
    "    [0.7, 0.2, 0.1],\n",
    "    ])\n",
    "\n",
    "q = [\n",
    "    {0:0, 1:0, 2:0,},\n",
    "    {0:0, 1:0, 2:0,},\n",
    "    {0:0, 1:0, 2:0,},\n",
    "]\n",
    "yita = p_hit = 0.5\n",
    "min_prob = 0.2\n",
    "\n",
    "\n",
    "def random_process_with_clamp(probs, hit):\n",
    "    with torch.no_grad():\n",
    "        max_place = probs.argmax(-1, keepdims=True)\n",
    "        mask_max = torch.zeros_like(probs).scatter_(-1, max_place, 1).bool()\n",
    "        pmax = probs[mask_max].reshape(max_place.shape) #probs[max_place].clone()\n",
    "        if hit:\n",
    "            assert max_place.shape[-1] == 1\n",
    "            return max_place.squeeze(-1)\n",
    "        else:\n",
    "            # forbit max prob being chosen\n",
    "            # pmax = probs.max(axis=-1) #probs[max_place].clone()\n",
    "            yita_arr = torch.ones_like(pmax)*yita\n",
    "            yita_arr_clip = torch.minimum(pmax, yita_arr)\n",
    "            # p_hat = pmax + (pmax-1) / (1/yita_arr_clip-1) + 1e-10\n",
    "            p_hat = (pmax-yita_arr_clip)/(1-yita_arr_clip)\n",
    "            k = 1/(1-yita_arr_clip)\n",
    "            probs *= k\n",
    "            probs[mask_max] = p_hat.reshape(-1)\n",
    "\n",
    "            # print(probs)\n",
    "            dist = Categorical(probs=probs)\n",
    "            samp = dist.sample()\n",
    "            assert samp.shape[-1] != 1\n",
    "            return samp #.squeeze(-1)\n",
    "\n",
    "def random_process_with_clamp2(probs, hit):\n",
    "    with torch.no_grad():\n",
    "        max_place = probs.argmax(-1, keepdims=True)\n",
    "        mask_max = torch.zeros_like(probs).scatter_(-1, max_place, 1).bool()\n",
    "        pmax = probs[mask_max].reshape(max_place.shape) #probs[max_place].clone()\n",
    "        if hit:\n",
    "            assert max_place.shape[-1] == 1\n",
    "            return max_place.squeeze(-1)\n",
    "        else:\n",
    "            # forbit max prob being chosen\n",
    "            # pmax = probs.max(axis=-1) #probs[max_place].clone()\n",
    "            yita_arr = torch.ones_like(pmax)*yita\n",
    "            # p_hat = pmax + (pmax-1) / (1/yita_arr_clip-1) + 1e-10\n",
    "            p_hat = (pmax-yita_arr)/(1-yita_arr)\n",
    "            p_hat = p_hat.clamp(min=min_prob)\n",
    "            k = (1-p_hat)/(1-pmax)\n",
    "            probs *= k\n",
    "            probs[mask_max] = p_hat.reshape(-1)\n",
    "\n",
    "            # print(probs)\n",
    "            dist = Categorical(probs=probs)\n",
    "            samp = dist.sample()\n",
    "            assert samp.shape[-1] != 1\n",
    "            return samp #.squeeze(-1)\n",
    "\n",
    "# [0.3, 0.3, 0.4],\n",
    "# [0.2, 0.5, 0.3],\n",
    "# [0.7, 0.2, 0.1],\n",
    "\n",
    "# tensor([[0.4500, 0.4500, 0.1000],\n",
    "#         [0.3600, 0.1000, 0.5400],\n",
    "#         [0.4000, 0.4000, 0.2000]])\n",
    "for i in range(10000):\n",
    "    hit = True if torch.rand(()) < p_hit else False\n",
    "    res = random_process_with_clamp(probs.clone(), hit)\n",
    "    res2 = random_process_with_clamp2(probs.clone(), hit)\n",
    "    hit = True if torch.rand(()) < p_hit else False\n",
    "    for j in range(3):\n",
    "        q[j][res[j].item()] += 1\n",
    "print(q)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np \n",
    "n_agents = 10\n",
    "seed = 0.4\n",
    "probs_list = [0.1,0.2,0.7]\n",
    "reward_which_act_list = [5,6,7]\n",
    "def _schochestic_reward_(act, seed, probs_list, reward_which_act_list):\n",
    "    s = 0\n",
    "    ratio = 2\n",
    "    for i in range(len(probs_list)):\n",
    "        s += probs_list[i]\n",
    "        if seed <= s:\n",
    "            reward_which_act = reward_which_act_list[i]\n",
    "            break\n",
    "        \n",
    "    return sum(act==reward_which_act)/(n_agents)*ratio\n",
    "\n",
    "act = np.array([7,7,7,7,7,7,7,7,7,7])\n",
    "act = np.array([7,7,7,7,7,7,7,7,7,7])*0 + 8\n",
    "_schochestic_reward_(act, seed, probs_list, reward_which_act_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.0216)\n",
      "tensor(0.0226)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch.distributions.categorical import Categorical\n",
    "from torch.nn.functional import kl_div\n",
    "import torch.nn.functional as F\n",
    "\n",
    "probs = torch.Tensor(\n",
    "    [0.4, 0.6])\n",
    "probs2 = torch.Tensor(\n",
    "    [0.3, 0.7])\n",
    "\n",
    "print(kl_div(probs.log(), probs2, reduction='sum'))\n",
    "print(kl_div(probs2.log(), probs, reduction='sum'))\n",
    "# print(kl_div(probs2.log(), probs, reduction='sum'))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "client.send_targeted_dgram('ddd')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.5507)"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch.nn.functional as F\n",
    "x = torch.Tensor(\n",
    "    [0.4, 0.6])\n",
    "y = torch.Tensor(\n",
    "    [0.9, 0.1])\n",
    "kl = F.kl_div(x.log(), y, reduction='sum')\n",
    "kl  # kl 散度越大，说明概率分布的差异越大"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.7981, -0.5981])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.softmax(dim=-1).log()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.0000, 0.0226, 0.0226],\n",
       "        [0.0216, 0.0000, 0.0000],\n",
       "        [0.0216, 0.0000, 0.0000]])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "from torch.distributions.categorical import Categorical\n",
    "from torch.nn.functional import kl_div\n",
    "import torch.nn.functional as F\n",
    "from UTILS.tensor_ops import repeat_at\n",
    "probs = torch.Tensor(\n",
    "    [\n",
    "        [0.4, 0.6],\n",
    "        [0.3, 0.7],\n",
    "        [0.3, 0.7]\n",
    "    ]\n",
    "    )\n",
    "\n",
    "\n",
    "# probs.S # (?, n_agent, n_action)\n",
    "n_agent = probs.shape[-2]\n",
    "probs_rep = repeat_at(tensor=probs, insert_dim=-2, n_times=n_agent)\n",
    "# probs_rep.S # (?, n_agent, n_agent, n_action)\n",
    "probs_rep_transpose = probs_rep.swapaxes(-2,-3)\n",
    "mat = (probs_rep*probs_rep.log()-probs_rep*probs_rep_transpose.log()).sum(-1)\n",
    "mat # (?, n_agent, n_agent)\n",
    "# F.kl_div(probs_rep.log(), probs_rep_transpose, reduction='batchmean')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[0. 2. 1. 1. 2. 1. 1. 2. 2. 2. 0. 2. 1. 0. 2. 0. 2. 2.]\n",
      " [2. 0. 2. 2. 1. 2. 2. 0. 0. 1. 2. 1. 2. 2. 0. 2. 1. 1.]\n",
      " [1. 2. 0. 0. 2. 0. 0. 2. 2. 2. 1. 2. 0. 1. 2. 1. 2. 2.]\n",
      " [1. 2. 0. 0. 2. 0. 0. 2. 2. 2. 1. 2. 0. 1. 2. 1. 2. 2.]\n",
      " [2. 1. 2. 2. 0. 2. 2. 1. 1. 0. 2. 0. 2. 2. 1. 2. 0. 0.]\n",
      " [1. 2. 0. 0. 2. 0. 0. 2. 2. 2. 1. 2. 0. 1. 2. 1. 2. 2.]\n",
      " [1. 2. 0. 0. 2. 0. 0. 2. 2. 2. 1. 2. 0. 1. 2. 1. 2. 2.]\n",
      " [2. 0. 2. 2. 1. 2. 2. 0. 0. 1. 2. 1. 2. 2. 0. 2. 1. 1.]\n",
      " [2. 0. 2. 2. 1. 2. 2. 0. 0. 1. 2. 1. 2. 2. 0. 2. 1. 1.]\n",
      " [2. 1. 2. 2. 0. 2. 2. 1. 1. 0. 2. 0. 2. 2. 1. 2. 0. 0.]\n",
      " [0. 2. 1. 1. 2. 1. 1. 2. 2. 2. 0. 2. 1. 0. 2. 0. 2. 2.]\n",
      " [2. 1. 2. 2. 0. 2. 2. 1. 1. 0. 2. 0. 2. 2. 1. 2. 0. 0.]\n",
      " [1. 2. 0. 0. 2. 0. 0. 2. 2. 2. 1. 2. 0. 1. 2. 1. 2. 2.]\n",
      " [0. 2. 1. 1. 2. 1. 1. 2. 2. 2. 0. 2. 1. 0. 2. 0. 2. 2.]\n",
      " [2. 0. 2. 2. 1. 2. 2. 0. 0. 1. 2. 1. 2. 2. 0. 2. 1. 1.]\n",
      " [0. 2. 1. 1. 2. 1. 1. 2. 2. 2. 0. 2. 1. 0. 2. 0. 2. 2.]\n",
      " [2. 1. 2. 2. 0. 2. 2. 1. 1. 0. 2. 0. 2. 2. 1. 2. 0. 0.]\n",
      " [2. 1. 2. 2. 0. 2. 2. 1. 1. 0. 2. 0. 2. 2. 1. 2. 0. 0.]]\n",
      "[[ 0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      " [ 1  0  1  1  0  1  1  0  0  0  1  0  1  1  0  1  0  0]\n",
      " [ 3  2  1  1  0  1  1  2  2  0  3  0  1  3  2  3  0  0]\n",
      " [ 3  5  6  6  4  1  1  5  2  0  7  0  1  7  2  3  4  0]\n",
      " [ 3 11  6 13  9  1  1  5 10  0  7  0 12 15  2 14  4  8]\n",
      " [ 3 11  6 13  9  1 17  5 10 16  7  0 12 15  2 14  4  8]]\n"
     ]
    }
   ],
   "source": [
    "# n_agent = 16\n",
    "tree = get_division_tree(18)\n",
    "# current_level = 3\n",
    "# blood_distance = np.ones(shape=(n_agent,n_agent), dtype=np.float64) * np.nan\n",
    "# for i in range(n_agent):\n",
    "#     for j in range(n_agent):\n",
    "#         if i==j:blood_distance[i,j] = 0\n",
    "#         if not np.isnan(blood_distance[i,j]):\n",
    "#             continue\n",
    "#         for t in range(current_level+1):\n",
    "#             investigate_level = (current_level) - t\n",
    "#             if tree[investigate_level, i] == tree[investigate_level, j]:\n",
    "#                 blood_distance[i,j] = t\n",
    "#                 blood_distance[j,i] = t\n",
    "#                 break\n",
    "\n",
    "\n",
    "blood_distance = get_blood_distance(tree, 2)\n",
    "print(blood_distance)\n",
    "print(tree)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],\n",
       "       [ 1,  0,  1,  0,  1,  1,  1,  1,  0,  1,  0,  0,  0,  1,  0,  0],\n",
       "       [ 1,  0,  3,  2,  3,  1,  1,  1,  2,  3,  0,  2,  2,  3,  0,  0],\n",
       "       [ 1,  0,  3,  5,  7,  6,  1,  6,  2,  7,  4,  5,  2,  3,  4,  0],\n",
       "       [12,  0,  3, 11, 15,  6,  1, 13, 10,  7,  9,  5,  2, 14,  4,  8]])"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "[[ 0.  2.  2. nan  2. nan  2. nan nan  2. nan  2. nan nan  2. nan]\n",
    " [ 2.  0.  2. nan  2. nan  2. nan nan  2. nan  2. nan nan  2. nan]\n",
    " [ 2.  2.  0. nan  2. nan  2. nan nan  2. nan  2. nan nan  2. nan]\n",
    " [nan nan nan  0. nan  2. nan  2.  2. nan  2. nan  2.  2. nan  2.]\n",
    " [ 2.  2.  2. nan  0. nan  2. nan nan  2. nan  2. nan nan  2. nan]\n",
    " [nan nan nan  2. nan  0. nan  2.  2. nan  2. nan  2.  2. nan  2.]\n",
    " [ 2.  2.  2. nan  2. nan  0. nan nan  2. nan  2. nan nan  2. nan]\n",
    " [nan nan nan  2. nan  2. nan  0.  2. nan  2. nan  2.  2. nan  2.]\n",
    " [nan nan nan  2. nan  2. nan  2.  0. nan  2. nan  2.  2. nan  2.]\n",
    " [ 2.  2.  2. nan  2. nan  2. nan nan  0. nan  2. nan nan  2. nan]\n",
    " [nan nan nan  2. nan  2. nan  2.  2. nan  0. nan  2.  2. nan  2.]\n",
    " [ 2.  2.  2. nan  2. nan  2. nan nan  2. nan  0. nan nan  2. nan]\n",
    " [nan nan nan  2. nan  2. nan  2.  2. nan  2. nan  0.  2. nan  2.]\n",
    " [nan nan nan  2. nan  2. nan  2.  2. nan  2. nan  2.  0. nan  2.]\n",
    " [ 2.  2.  2. nan  2. nan  2. nan nan  2. nan  2. nan nan  0. nan]\n",
    " [nan nan nan  2. nan  2. nan  2.  2. nan  2. nan  2.  2. nan  0.]]\n",
    "\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'d1': tensor([1.1618e-13, 4.3868e-16, 3.4678e-18, 0.0000e+00, 1.0000e+00, 8.3701e-14,\n",
       "         1.8273e-12], device='cuda:0', requires_grad=True),\n",
       " 'd2': tensor([0.0756, 0.2109, 0.0628, 0.0104, 0.3506, 0.1273, 0.1624],\n",
       "        device='cuda:0', requires_grad=True)}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pickle\n",
    "with open('RECYCLE/wifi', 'rb') as f:\n",
    "    p_list = pickle.load(f)\n",
    "\n",
    "p_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([5.0909e-02, 1.9223e-04, 1.5197e-06, 8.2279e-11, 1.1147e-01, 3.6681e-02,\n",
      "        8.0075e-01], dtype=torch.float64)\n",
      "tensor([0.0756, 0.2109, 0.0628, 0.0104, 0.3506, 0.1273, 0.1624],\n",
      "       dtype=torch.float64)\n",
      "tensor(1.0829, dtype=torch.float64)\n",
      "tensor(2.6698, dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn.functional as F\n",
    "from torch.distributions.categorical import Categorical\n",
    "from torch.nn.functional import kl_div\n",
    "\n",
    "\n",
    "logits1 = torch.Tensor([  4.4786,  -1.1005,  -5.9407, -15.7646,  5.2623,   4.1508,   7.2341]).type(torch.DoubleTensor)\n",
    "logits2 = torch.Tensor([ -0.5480,  0.4779, -0.7330, -2.5290,  0.9862, -0.0273,  0.2163]).type(torch.DoubleTensor)\n",
    "\n",
    "dist1 = Categorical(logits = logits1)\n",
    "dist2 = Categorical(logits = logits2)\n",
    "print(dist1.probs)\n",
    "print(dist2.probs)\n",
    "from torch.distributions import kl_divergence\n",
    "\n",
    "print(kl_divergence(dist1, dist2))\n",
    "print(kl_divergence(dist2, dist1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1.1618e-13, 4.3868e-16, 3.4678e-18, 0.0000e+00, 1.0000e+00, 8.3701e-14,\n",
       "        1.8273e-12], device='cuda:0', requires_grad=True)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# probs = torch.Tensor(\n",
    "#     [0, 1])\n",
    "# probs2 = torch.Tensor(\n",
    "#     [0.3, 0.7])\n",
    "\n",
    "# probs = p_list['d1']\n",
    "# probs2 = p_list['d2']\n",
    "\n",
    "# print(kl_div(probs.log(), probs2))\n",
    "# print(kl_div(probs2.log(), probs))\n",
    "\n",
    "# probs = torch.Tensor(\n",
    "#     [0, 1])\n",
    "# probs2 = torch.Tensor(\n",
    "#     [0.3, 0.7])\n",
    "\n",
    "# probs = p_list['d1']\n",
    "# probs2 = p_list['d2']\n",
    "\n",
    "# print(kl_div(probs.log(), probs2))\n",
    "# print(kl_div(probs2.log(), probs))\n",
    "\n",
    "# probs = torch.Tensor(\n",
    "#     [0, 1])\n",
    "# probs2 = torch.Tensor(\n",
    "#     [0.3, 0.7])\n",
    "\n",
    "# probs = p_list['d1']\n",
    "# probs2 = p_list['d2']\n",
    "\n",
    "# print(kl_div(probs.log(), probs2))\n",
    "# print(kl_div(probs2.log(), probs))\n",
    "\n",
    "# probs = torch.Tensor(\n",
    "#     [0, 1])\n",
    "# probs2 = torch.Tensor(\n",
    "#     [0.3, 0.7])\n",
    "\n",
    "# probs = p_list['d1']\n",
    "# probs2 = p_list['d2']\n",
    "\n",
    "# print(kl_div(probs.log(), probs2))\n",
    "# print(kl_div(probs2.log(), probs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[ 0.,  1.,  2.],\n",
      "         [ 3.,  4.,  5.]],\n",
      "\n",
      "        [[ 6.,  7.,  8.],\n",
      "         [ 9., 10., 11.]],\n",
      "\n",
      "        [[12., 13., 14.],\n",
      "         [15., 16., 17.]]])\n",
      "tensor([[[ 0.,  0.,  0.],\n",
      "         [ 3.,  4.,  5.]],\n",
      "\n",
      "        [[ 6.,  7.,  8.],\n",
      "         [ 0.,  0.,  0.]],\n",
      "\n",
      "        [[ 0.,  0.,  0.],\n",
      "         [15., 16., 17.]]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from UTILS.tensor_ops import gather_righthand, repeat_at\n",
    "def scatter_righthand(scatter_into, src, index, check=True):\n",
    "    index = index.long()\n",
    "    i_dim = index.dim()\n",
    "    s_dim = src.dim()\n",
    "    t_dim = i_dim - 1\n",
    "    index_new_shape = list(src.shape)\n",
    "    index_new_shape[t_dim] = index.shape[t_dim]\n",
    "    for _ in range(i_dim, s_dim):\n",
    "        index = index.unsqueeze(-1)\n",
    "    index_expand = index.expand(index_new_shape)  # only this two line matters\n",
    "    return scatter_into.scatter(t_dim, index_expand, src)\n",
    "\n",
    "orig = torch.Tensor([[[ 0,  1,  2], [ 3,  4,  5]],\n",
    "                    [[ 6,  7,  8], [ 9, 10, 11]],\n",
    "                    [[12, 13, 14], [15, 16, 17]]])\n",
    "index = torch.Tensor([[0], [1], [0]])\n",
    "print(orig)\n",
    "\n",
    "res      = gather_righthand(src=orig, index=index)\n",
    "res[:] = 0\n",
    "\n",
    "orig_fix = orig.clone().detach()\n",
    "orig_fix = scatter_righthand(orig_fix, src=res, index=index)\n",
    "print(orig_fix)\n",
    "\n",
    "\n",
    "# orig.scatter(dim=1,index=index.long(),src=res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "s1\n",
      "s1\n",
      "s2\n",
      "s2\n",
      "torch.Size([64, 16, 8, 88, 888])\n"
     ]
    }
   ],
   "source": [
    "index = torch.randint(high=16,size=(64,5))\n",
    "print('s1')\n",
    "res = gather_righthand(src,index)\n",
    "print('s1')\n",
    "res[:] = 0\n",
    "print('s2')\n",
    "resX = scatter_righthand(scatter_into=src, src=res, index=index)\n",
    "print('s2')\n",
    "print(resX.S)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x2 tensor([[101., 101., 101., 101., 100., 100., 100.],\n",
      "        [202., 202., 202., 200., 201., 200., 200.],\n",
      "        [303., 303., 303., 300., 300., 301., 300.],\n",
      "        [404., 404., 404., 400., 400., 400., 401.]])\n",
      "perm_index tensor([2, 1, 3, 0])\n",
      "confact_x2 tensor([[401., 401., 401., 400., 400., 400., 401.],\n",
      "        [202., 202., 202., 200., 201., 200., 200.],\n",
      "        [103., 103., 103., 101., 100., 100., 100.],\n",
      "        [304., 304., 304., 300., 300., 301., 300.]])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from UTILS.tensor_ops import gather_righthand, repeat_at\n",
    "from UTILS.tensor_ops import add_onehot_id_at_last_dim, repeat_at, _2tensor, gather_righthand, scatter_righthand\n",
    "x_in = torch.Tensor([\n",
    "    [1,1,1],\n",
    "    [2,2,2],\n",
    "    [3,3,3],\n",
    "    [4,4,4],\n",
    "])\n",
    "n_agent = 4\n",
    "\n",
    "\n",
    "nets = [\n",
    "    lambda x: x+100,\n",
    "    lambda x: x+200,\n",
    "    lambda x: x+300,\n",
    "    lambda x: x+400\n",
    "]\n",
    "\n",
    "\n",
    "x0 = add_onehot_id_at_last_dim(x_in)\n",
    "# x1 = self.shared_net(x0)\n",
    "res = []\n",
    "for i in range(n_agent):\n",
    "    res.append(nets[i](x0[..., i, :]))\n",
    "x2 = torch.stack(res, -2)\n",
    "# x22 = self.nets[0](x1)\n",
    "print('x2',x2)\n",
    "\n",
    "######### forward twice: shuffle forward\n",
    "perm_index = torch.randperm(n_agent, device=x_in.device)   # shape = (n_agent)\n",
    "print('perm_index',perm_index)\n",
    "\n",
    "perm_index = perm_index.expand(x_in.shape[:-1]) # shape = (...?, n_agent)\n",
    "# x_in shape = (...?, n_agent, coredim)\n",
    "perm_x_in = gather_righthand(src=x_in, index=perm_index, check=False)\n",
    "perm_x0 = add_onehot_id_at_last_dim(perm_x_in)\n",
    "perm_res = []\n",
    "for i in range(n_agent):\n",
    "    perm_res.append(nets[i](perm_x0[..., i, :]))\n",
    "perm_x2 = torch.stack(perm_res, -2)\n",
    "\n",
    "# 103\n",
    "# 202\n",
    "# 304\n",
    "# 401\n",
    "\n",
    "\n",
    "# 401 vs 101\n",
    "# 202 vs 202\n",
    "# 103 vs 303\n",
    "# 304 vs 404\n",
    "\n",
    "confact_x2 = torch.zeros_like(perm_x2) + np.nan\n",
    "confact_x2 = scatter_righthand(scatter_into=confact_x2, src=perm_x2, index=perm_index)\n",
    "print('confact_x2',confact_x2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([False, False, False, False,  True])\n",
      "tensor([ True, False,  True, False, False])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "res=torch.isnan(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))\n",
    "print(res)\n",
    "\n",
    "res=torch.isfinite(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))\n",
    "print(res)\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}