{
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {
            "display_name": "Python",
            "language": "python",
            "name": "python"
        },
        "language_info": {
            "name": "python"
        }
    },
    "cells": [
        {
            "id": "1df15b1b",
            "cell_type": "markdown",
            "source": "<div class='alert alert-warning'>\n\nSciPy's interactive examples with Jupyterlite are experimental and may not always work as expected. Execution of cells containing imports may result in large downloads (up to 60MB of content for the first import from SciPy). Load times when importing from SciPy may take roughly 10-20 seconds. If you notice any problems, feel free to open an [issue](https://github.com/scipy/scipy/issues/new/choose).\n\n</div>",
            "metadata": {}
        },
        {
            "id": "a4ca5295",
            "cell_type": "markdown",
            "source": "Here we reproduce a number of the numerical examples from [1] as a demonstration\nof the functionality offered by this method.\n",
            "metadata": {}
        },
        {
            "id": "7f3baeb8",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "import numpy as np\nimport matplotlib.pyplot as plt\nfrom scipy.interpolate import AAA\nimport warnings",
            "outputs": []
        },
        {
            "id": "4ffc066d",
            "cell_type": "markdown",
            "source": "For the first example we approximate the gamma function on ``[-3.5, 4.5]`` by\nextrapolating from 100 samples in ``[-1.5, 1.5]``.\n",
            "metadata": {}
        },
        {
            "id": "b25210e6",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "from scipy.special import gamma\nsample_points = np.linspace(-1.5, 1.5, num=100)\nr = AAA(sample_points, gamma(sample_points))\nz = np.linspace(-3.5, 4.5, num=1000)\nfig, ax = plt.subplots()\nax.plot(z, gamma(z), label=\"Gamma\")\nax.plot(sample_points, gamma(sample_points), label=\"Sample points\")\nax.plot(z, r(z).real, '--', label=\"AAA approximation\")\nax.set(xlabel=\"z\", ylabel=\"r(z)\", ylim=[-8, 8], xlim=[-3.5, 4.5])\nax.legend()\nplt.show()",
            "outputs": []
        },
        {
            "id": "8d44d64e",
            "cell_type": "markdown",
            "source": "We can also view the poles of the rational approximation and their residues:\n",
            "metadata": {}
        },
        {
            "id": "11bbef2b",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "order = np.argsort(r.poles())\nr.poles()[order]",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([-3.81591039e+00+0.j        , -3.00269049e+00+0.j        ,\n       -1.99999988e+00+0.j        , -1.00000000e+00+0.j        ,\n        5.85842812e-17+0.j        ,  4.77485458e+00-3.06919376j,\n        4.77485458e+00+3.06919376j,  5.29095868e+00-0.97373072j,\n        5.29095868e+00+0.97373072j])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "1d8264fb",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "r.residues()[order]",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([ 0.03658074 +0.j        , -0.16915426 -0.j        ,\n        0.49999915 +0.j        , -1.         +0.j        ,\n        1.         +0.j        , -0.81132013 -2.30193429j,\n       -0.81132013 +2.30193429j,  0.87326839+10.70148546j,\n        0.87326839-10.70148546j])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "d9d611d2",
            "cell_type": "markdown",
            "source": "For the second example, we call `AAA` with a spiral of 1000 points that wind 7.5\ntimes around the origin in the complex plane.\n",
            "metadata": {}
        },
        {
            "id": "92d3bf8b",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "z = np.exp(np.linspace(-0.5, 0.5 + 15j*np.pi, 1000))\nr = AAA(z, np.tan(np.pi*z/2), rtol=1e-13)",
            "outputs": []
        },
        {
            "id": "5579bcff",
            "cell_type": "markdown",
            "source": "We see that AAA takes 12 steps to converge with the following errors:\n",
            "metadata": {}
        },
        {
            "id": "dfc7b476",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "r.errors.size",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "12"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "05edfd48",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "r.errors",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([2.49261500e+01, 4.28045609e+01, 1.71346935e+01, 8.65055336e-02,\n       1.27106444e-02, 9.90889874e-04, 5.86910543e-05, 1.28735561e-06,\n       3.57007424e-08, 6.37007837e-10, 1.67103357e-11, 1.17112299e-13])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "29489ac9",
            "cell_type": "markdown",
            "source": "We can also plot the computed poles:\n",
            "metadata": {}
        },
        {
            "id": "352a473d",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "fig, ax = plt.subplots()\nax.plot(z.real, z.imag, '.', markersize=2, label=\"Sample points\")\nax.plot(r.poles().real, r.poles().imag, '.', markersize=5,\n        label=\"Computed poles\")\nax.set(xlim=[-3.5, 3.5], ylim=[-3.5, 3.5], aspect=\"equal\")\nax.legend()\nplt.show()",
            "outputs": []
        },
        {
            "id": "95bbbf76",
            "cell_type": "markdown",
            "source": "We now demonstrate the removal of Froissart doublets using the `clean_up` method\nusing an example from [1]. Here we approximate the function\n$f(z)=\\log(2 + z^4)/(1 + 16z^4)$ by sampling it at 1000 roots of unity. The\nalgorithm is run with ``rtol=0`` and ``clean_up=False`` to deliberately cause\nFroissart doublets to appear.\n",
            "metadata": {}
        },
        {
            "id": "6f6222c4",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "z = np.exp(1j*2*np.pi*np.linspace(0,1, num=1000))\ndef f(z):\n    return np.log(2 + z**4)/(1 - 16*z**4)\nwith warnings.catch_warnings():  # filter convergence warning due to rtol=0\n    warnings.simplefilter('ignore', RuntimeWarning)\n    r = AAA(z, f(z), rtol=0, max_terms=50, clean_up=False)\nmask = np.abs(r.residues()) < 1e-13\nfig, axs = plt.subplots(ncols=2)\naxs[0].plot(r.poles().real[~mask], r.poles().imag[~mask], '.')\naxs[0].plot(r.poles().real[mask], r.poles().imag[mask], 'r.')",
            "outputs": []
        },
        {
            "id": "84408531",
            "cell_type": "markdown",
            "source": "Now we call the `clean_up` method to remove Froissart doublets.\n",
            "metadata": {}
        },
        {
            "id": "d5e5e2ce",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "with warnings.catch_warnings():\n    warnings.simplefilter('ignore', RuntimeWarning)\n    r.clean_up()",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "4  # may vary"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "370b415c",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "mask = np.abs(r.residues()) < 1e-13\naxs[1].plot(r.poles().real[~mask], r.poles().imag[~mask], '.')\naxs[1].plot(r.poles().real[mask], r.poles().imag[mask], 'r.')\nplt.show()",
            "outputs": []
        },
        {
            "id": "845ad775",
            "cell_type": "markdown",
            "source": "The left image shows the poles prior of the approximation ``clean_up=False`` with\npoles with residue less than ``10^-13`` in absolute value shown in red. The right\nimage then shows the poles after the `clean_up` method has been called.\n",
            "metadata": {}
        }
    ]
}