{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Diagonal\nThis example shows how to use the :py:class:`pylops_distributed.Diagonal`\noperator to perform *Element-wise multiplication* between the input\nvector and a vector $\\mathbf{d}$.\n\nIn other words, the operator acts as a  diagonal operator $\\mathbf{D}$\nwhose elements along the diagonal are the elements of the vector\n$\\mathbf{d}$.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as pltgs\nimport dask.array as da\n\nimport pylops_distributed\n\nplt.close('all')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's define a diagonal operator $\\mathbf{d}$ with increasing numbers\nfrom ``0`` to ``N`` and a unitary model $\\mathbf{x}$.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "N = 10\nd = da.arange(N, chunks=N//2)\nx = da.ones(N, chunks=N//2)\n\nDop = pylops_distributed.Diagonal(d, compute=(True, True))\n\ny = Dop * x\ny1 = Dop.H * x\n\ngs = pltgs.GridSpec(1, 6)\nfig = plt.figure(figsize=(7, 3))\nax = plt.subplot(gs[0, 0:3])\nim = ax.imshow(Dop.Op.matrix(), cmap='rainbow', vmin=0, vmax=N)\nax.set_title('A', size=20, fontweight='bold')\nax.set_xticks(np.arange(N-1)+0.5)\nax.set_yticks(np.arange(N-1)+0.5)\nax.grid(linewidth=3, color='white')\nax.xaxis.set_ticklabels([])\nax.yaxis.set_ticklabels([])\nax.axis('tight')\nax = plt.subplot(gs[0, 3])\nax.imshow(x[:, np.newaxis], cmap='rainbow', vmin=0, vmax=N)\nax.set_title('x', size=20, fontweight='bold')\nax.set_xticks([])\nax.set_yticks(np.arange(N-1)+0.5)\nax.grid(linewidth=3, color='white')\nax.xaxis.set_ticklabels([])\nax.yaxis.set_ticklabels([])\nax = plt.subplot(gs[0, 4])\nax.text(0.35, 0.5, '=', horizontalalignment='center',\n        verticalalignment='center', size=40, fontweight='bold')\nax.axis('off')\nax = plt.subplot(gs[0, 5])\nax.imshow(y[:, np.newaxis], cmap='rainbow', vmin=0, vmax=N)\nax.set_title('y', size=20, fontweight='bold')\nax.set_xticks([])\nax.set_yticks(np.arange(N - 1) + 0.5)\nax.grid(linewidth=3, color='white')\nax.xaxis.set_ticklabels([])\nax.yaxis.set_ticklabels([])\nfig.colorbar(im, ax=ax, ticks=[0, N], pad=0.3, shrink=0.7)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Similarly we can consider the input model as composed of two or more\ndimensions. In this case the diagonal operator can be still applied to\neach element or broadcasted along a specific direction. Let's start with the\nsimplest case where each element is multipled by a different value\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nx, ny = 6, 5\nx = da.ones((nx, ny),  chunks=(nx//2, ny))\nprint('x =\\n%s' % x)\n\nd = da.arange(nx*ny, chunks=((ny*nx)//2)).reshape(nx, ny)\nDop = pylops_distributed.Diagonal(d, compute=(True, True))\n\ny = Dop * x.flatten()\ny1 = Dop.H * x.flatten()\n\nprint('y = D*x =\\n%s' % y.reshape(nx, ny))\nprint('xadj = D\\'*x =\\n%s ' % y1.reshape(nx, ny))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "And we now broadcast\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "nx, ny = 6, 5\nx = da.ones((nx, ny),  chunks=(nx//2, ny))\nprint('x =\\n%s' % x)\n\n# 1st dim\nd = da.arange(nx, chunks=nx//2)\nDop = pylops_distributed.Diagonal(d, dims=(nx, ny), dir=0)\n\ny = Dop * x.flatten()\ny1 = Dop.H * x.flatten()\n\nprint('1st dim: y = D*x =\\n%s' % y.reshape(nx, ny))\nprint('1st dim: xadj = D\\'*x =\\n%s ' % y1.reshape(nx, ny))\n\n# 2nd dim\nd = da.arange(ny, chunks=nx//2)\nDop = pylops_distributed.Diagonal(d, dims=(nx, ny), dir=1)\n\ny = Dop * x.flatten()\ny1 = Dop.H * x.flatten()\n\nprint('2nd dim: y = D*x =\\n%s' % y.reshape(nx, ny))\nprint('2nd dim: xadj = D\\'*x =\\n%s ' % y1.reshape(nx, ny))"
      ]
    }
  ],
  "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.6.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}