{
  "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.Diagonal` operator\nto perform *Element-wise multiplication* between the input vector and a vector $\\mathbf{d}$.\n\nIn other words, the operator acts as a  diagonal operator $\\mathbf{D}$ whose elements along\nthe diagonal are the elements of the vector $\\mathbf{d}$.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import matplotlib.gridspec as pltgs\nimport matplotlib.pyplot as plt\nimport numpy as np\n\nimport pylops\n\nplt.close(\"all\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's define a diagonal operator $\\mathbf{d}$ with increasing numbers from\n``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 = np.arange(N)\nx = np.ones(N)\n\nDop = pylops.Diagonal(d)\n\ny = Dop * x\ny1 = Dop.H * x\n\ngs = pltgs.GridSpec(1, 6)\nfig = plt.figure(figsize=(7, 4))\nax = plt.subplot(gs[0, 0:3])\nim = ax.imshow(Dop.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(\n    0.35,\n    0.5,\n    \"=\",\n    horizontalalignment=\"center\",\n    verticalalignment=\"center\",\n    size=40,\n    fontweight=\"bold\",\n)\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)\nplt.tight_layout()"
      ]
    },
    {
      "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 = 3, 5\nx = np.ones((nx, ny))\nprint(f\"x =\\n{x}\")\n\nd = np.arange(nx * ny).reshape(nx, ny)\nDop = pylops.Diagonal(d)\n\ny = Dop * x.ravel()\ny1 = Dop.H * x.ravel()\n\nprint(f\"y = D*x =\\n{y.reshape(nx, ny)}\")\nprint(f\"xadj = D'*x =\\n{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 = 3, 5\nx = np.ones((nx, ny))\nprint(f\"x =\\n{x}\")\n\n# 1st dim\nd = np.arange(nx)\nDop = pylops.Diagonal(d, dims=(nx, ny), axis=0)\n\ny = Dop * x.ravel()\ny1 = Dop.H * x.ravel()\n\nprint(f\"1st dim: y = D*x =\\n{y.reshape(nx, ny)}\")\nprint(f\"1st dim: xadj = D'*x =\\n{y1.reshape(nx, ny)}\")\n\n# 2nd dim\nd = np.arange(ny)\nDop = pylops.Diagonal(d, dims=(nx, ny), axis=1)\n\ny = Dop * x.ravel()\ny1 = Dop.H * x.ravel()\n\nprint(f\"2nd dim: y = D*x =\\n{y.reshape(nx, ny)}\")\nprint(f\"2nd dim: xadj = D'*x =\\n{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.9.15"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}