SO73466464.ipynb

so73466464.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SO73466464.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyO6Ga9gf7HRbPS90OBX3ph3",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/gist/StefRe/a3080dd8998c749a914f879e48ad2ace/so73466464.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "l8RI8tNraM1V"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import numba\n",
        "\n",
        "@numba.njit\n",
        "def randperm(n):\n",
        "  \"\"\"Permuation of 1, 2, ... n and its inverse\"\"\"\n",
        "  p = np.arange(1, n+1, dtype=np.int32) # initialize with identity permutation\n",
        "  pinv = np.empty_like(p)\n",
        "  for i in range(n-1, 0, -1):           # loop over all items except the first one\n",
        "    z = np.random.randint(0, i+1)       \n",
        "    temp = p[z]                         # swap numbers at i and z\n",
        "    p[z] = p[i]\n",
        "    p[i] = temp\n",
        "    pinv[temp-1] = i                    # pinv[p[z]-1] = i\n",
        "  pinv[p[0]-1] = 0  \n",
        "  return p, pinv\n",
        "\n",
        "randperm(10) # compile it\n",
        "\n",
        "def invert_permutation_numpy2(permutation):\n",
        "    inv = np.empty_like(permutation)\n",
        "    inv[permutation-1] = np.arange(len(inv), dtype=inv.dtype)\n",
        "    return inv"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n = 100_000_000"
      ],
      "metadata": {
        "id": "D6oRuWrVaPl-"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "p = np.random.permutation(np.arange(1, n+1, dtype=np.int32))\n",
        "pinv = np.argsort(p)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CWmAg8MEavw5",
        "outputId": "b01db56a-9253-49ae-e002-f0e65a585781"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "27.5 s ± 537 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%%timeit\n",
        "p = np.random.permutation(np.arange(1, n+1, dtype=np.int32))\n",
        "invert_permutation_numpy2(p)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PkeZnEw4aU2C",
        "outputId": "0d6938ce-6494-4d6d-ea28-ad40a7d2e17c"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "8.74 s ± 465 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%timeit p, pinv = randperm(n)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hsVzTc8PakSz",
        "outputId": "ad70511a-aa45-41be-fa45-3661b8e6643c"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "3.66 s ± 10.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import platform\n",
        "import psutil\n",
        "\n",
        "if platform.system() == \"Windows\":\n",
        "    !wmic cpu get name | find \"@\"\n",
        "else:\n",
        "    !cat /proc/cpuinfo | grep \"model name\"\n",
        "\n",
        "print(f'{psutil.virtual_memory().total / 1024**3:.0f} GB')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "pWXQGXFba2OZ",
        "outputId": "97b9e4e3-b73e-495f-bacb-1d4e3bb48c2e"
      },
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "model name\t: Intel(R) Xeon(R) CPU @ 2.00GHz\n",
            "model name\t: Intel(R) Xeon(R) CPU @ 2.00GHz\n",
            "13 GB\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "np.__version__, numba.__version__"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PNIJfkwVb4cY",
        "outputId": "11e2bd63-15b7-470e-cf74-18777b8bf8a5"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "('1.21.6', '0.56.0')"
            ]
          },
          "metadata": {},
          "execution_count": 15
        }
      ]
    }
  ]
}

Hardware	OS	numpy/argsort	numpy/invert_permutation_numpy2	randperm
EPYC 7B12 - 2.25 GHz - 13 GB	Linux	20 s	6 s	4 s
Xeon - 2.00 GHz - 13 GB	Linux	28 s	8 s	4 s
Corei5-8400T - 1.70 GHz - 8 GB	Windows	34 s	8 s	13 s