argpermutations.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "60788a7a-f955-40f2-8456-ffc33b262f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import awkward as ak\n",
    "import numpy as np\n",
    "import numba as nb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "951e38a8-54e2-4eb4-bcec-974fc3e19ff9",
   "metadata": {},
   "source": [
    "Generate permutations using Heap's Algorithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c6137333-9400-4f04-80d7-26e834c52d55",
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def _permutations_k(A, k):\n",
    "    c = np.zeros(k, dtype=np.uint64)\n",
    "    \n",
    "    yield A\n",
    "    i = 0\n",
    "    while i < k:\n",
    "        if c[i] < i:\n",
    "            if i % 2:\n",
    "                A[c[i]], A[i] = A[i], A[c[i]]\n",
    "            else:\n",
    "                A[0], A[i] = A[i], A[0]\n",
    "            yield A\n",
    "            c[i] += 1\n",
    "            i = 0\n",
    "        else:\n",
    "            c[i] = 0\n",
    "            i += 1\n",
    "            \n",
    "@nb.njit\n",
    "def _permutations(A):\n",
    "    return _permutations_k(A, len(A))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25fdaac1-7cfe-4fd7-9377-7ebfe32c8aa9",
   "metadata": {},
   "source": [
    "Generate Permuted Indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "0d3fea7f-6570-4036-b2cb-cf7acf79124d",
   "metadata": {},
   "outputs": [],
   "source": [
    "@nb.njit\n",
    "def argpermutations(array, builder):\n",
    "    for x in array:\n",
    "        i = np.arange(len(x))\n",
    "        builder.begin_list()\n",
    "        for p in _permutations(i):\n",
    "            builder.begin_list()\n",
    "            for j in p:\n",
    "                builder.integer(j)\n",
    "            builder.end_list()\n",
    "        builder.end_list()\n",
    "    return builder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c10a775e-2567-4c3c-9a15-b886f74d0115",
   "metadata": {},
   "source": [
    "Create some test data!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "1c9f8f1b-f407-417d-922b-845ce6a90c8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = ak.Array(\n",
    "    [\n",
    "        [\n",
    "            {\"id\": 0, \"charge\": 1},\n",
    "            {\"id\": 1, \"charge\": -1},\n",
    "            {\"id\": 2, \"charge\": -1},\n",
    "            {\"id\": 3, \"charge\": 1},\n",
    "        ],\n",
    "        [\n",
    "            {\"id\": 0, \"charge\": -1},\n",
    "            {\"id\": 1, \"charge\": 1},\n",
    "            {\"id\": 2, \"charge\": -1},\n",
    "            {\"id\": 3, \"charge\": 1},\n",
    "            {\"id\": 4, \"charge\": -1},\n",
    "            {\"id\": 5, \"charge\": 1},\n",
    "        ],\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfa0f230-47a9-47f9-8e1f-f5b264ebe1e0",
   "metadata": {},
   "source": [
    "Generate the indices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "a5e33de2-1384-4ced-bde3-2514a81cdce6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ix = argpermutations(x, ak.ArrayBuilder()).snapshot()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "935163fe-4f4d-4987-a108-ff487bb20b27",
   "metadata": {},
   "source": [
    "Now, broadcast the array `x` to have the same datashape as `ix`. We do this by adding a new dimension that corresponds to the \"proposals\" dimension in `ix`. Then, we use broadcasting to grow this dimension to match `ix`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "26076dd5-c59d-4a8b-8e65-825dfaa748a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_p, _ = ak.broadcast_arrays(x[..., np.newaxis, :], ix)\n",
    "x_p_ix = x_p[ix]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95d43c17-cebe-4ece-80f0-83577139467f",
   "metadata": {},
   "source": [
    "Let's take only an even number of proposals, by computing the largest even number of items per proposal, and masking only items with a local index below this position."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "55720aa1-9a84-4383-a0fe-0f2bf91631e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_even = (ak.num(x_p_ix, axis=-1)//2) * 2\n",
    "is_within_n = ak.local_index(x_p_ix) < n_even"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22d9d1b8-68aa-48b3-ba66-bb25a2395935",
   "metadata": {},
   "source": [
    "Take only these within-n items"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "c0c3ab0c-708a-4add-a86b-40bfeb1a6053",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_p_ix_n = x_p_ix[is_within_n]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a422565-3a7b-469c-8ef8-b7c5dda0a1a7",
   "metadata": {},
   "source": [
    "Build pairs the adjacent entries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "e9ec6d93-71fa-45a4-bf94-ab2d6dfd9130",
   "metadata": {},
   "outputs": [],
   "source": [
    "pair = ak.zip((x_p_ix_n[..., ::2], x_p_ix_n[..., 1::2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0001086-4ef9-4d93-8377-78542a188e6f",
   "metadata": {},
   "source": [
    "Take only sum-zero charge combinations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "id": "50344817-f9c1-4605-915e-761828aa083f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_cfg =  pair[(pair.slot0.charge + pair.slot1.charge) == 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "id": "435f9b61-fc43-4394-87d5-12eb968921b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1328"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ak.count(pair_cfg.slot0.id)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}