{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Operators with Multiprocessing\nThis example shows how perform a scalability test for one of PyLops operators\nthat uses ``multiprocessing`` to spawn multiple processes. Operators that\nsupport such feature are :class:`pylops.basicoperators.VStack`,\n:class:`pylops.basicoperators.HStack`, and\n:class:`pylops.basicoperators.BlockDiagonal`, and\n:class:`pylops.basicoperators.Block`.\n\nIn this example we will consider the BlockDiagonal operator which contains\n:class:`pylops.basicoperators.MatrixMult` operators along its main diagonal.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\nimport numpy as np\n\nimport pylops\n\nplt.close(\"all\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's start by creating N MatrixMult operators and the BlockDiag operator\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "N = 100\nNops = 32\nOps = [pylops.MatrixMult(np.random.normal(0.0, 1.0, (N, N))) for _ in range(Nops)]\n\nOp = pylops.BlockDiag(Ops, nproc=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can now perform a scalability test on the forward operation\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "workers = [2, 3, 4]\ncompute_times, speedup = pylops.utils.multiproc.scalability_test(\n    Op, np.ones(Op.shape[1]), workers=workers, forward=True\n)\nplt.figure(figsize=(12, 3))\nplt.plot(workers, speedup, \"ko-\")\nplt.xlabel(\"# Workers\")\nplt.ylabel(\"Speed Up\")\nplt.title(\"Forward scalability test\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And likewise on the adjoint operation\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "compute_times, speedup = pylops.utils.multiproc.scalability_test(\n    Op, np.ones(Op.shape[0]), workers=workers, forward=False\n)\nplt.figure(figsize=(12, 3))\nplt.plot(workers, speedup, \"ko-\")\nplt.xlabel(\"# Workers\")\nplt.ylabel(\"Speed Up\")\nplt.title(\"Adjoint scalability test\")\nplt.tight_layout()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Note that we have not tested here the case with 1 worker. In this specific\ncase, since the computations are very small, the overhead of spawning processes\nis actually dominating the time of computations and so computing the\nforward and adjoint operations with a single worker is more efficient. We\nhope that this example can serve as a basis to inspect the scalability of\nmultiprocessing-enabled operators and choose the best number of processes.\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}