{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial FFT 2D parallel (MPI)\n",
    "\n",
    "In this tutorial, we present how to use fluidfft to perform 2D fft in parallel."
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "As in sequential, the first step would normally to benchmark the different methods for your case as explained in the page :ref:`bench`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Because, we are doing this tutorial in parallel with jupyter and ipyparallel, we first need to create an ipyparallel client and create a direct view as explained [here](http://ipyparallel.readthedocs.io). We previously started an ipcluster with the command `ipcluster start -n 4 --engines=MPIEngineSetLauncher`. This is just a jupyter/ipython thing and it has nothing to do with fluidfft."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ipyparallel as ipp\n",
    "rc = ipp.Client()\n",
    "dview = rc[:]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Afterwards, we will execute all cells in parallel so we always need to add the magic command `%%px` (see [here](http://ipyparallel.readthedocs.io/en/latest/magics.html))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] Hello world! I'm rank 0/4\n",
      "[stdout:1] Hello world! I'm rank 1/4\n",
      "[stdout:2] Hello world! I'm rank 2/4\n",
      "[stdout:3] Hello world! I'm rank 3/4\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "from fluiddyn.util.mpi import rank, nb_proc\n",
    "print(\"Hello world! I'm rank {}/{}\".format(rank, nb_proc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then it is very similar as in sequential so we do not need to explain too much!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "import numpy as np\n",
    "from fluidfft.fft2d import methods_mpi\n",
    "from fluidfft import import_fft_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] ['fft2d.mpi_with_fftwmpi2d', 'fft2d.mpi_with_fftw1d']\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "if rank == 0:\n",
    "    print(methods_mpi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "cls = import_fft_class('fft2d.mpi_with_fftw1d')\n",
    "o = cls(48, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] 1\n",
      "[stdout:1] 1\n",
      "[stdout:2] 1\n",
      "[stdout:3] 1\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "_ = o.run_tests()\n",
    "print(_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] t_fft = 7.07e-05 s; t_ifft = 2.05e-05 s\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "times = o.run_benchs()\n",
    "if rank == 0:\n",
    "    print('t_fft = {} s; t_ifft = {} s'.format(*times))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] True\n",
      "[stdout:1] True\n",
      "[stdout:2] True\n",
      "[stdout:3] True\n"
     ]
    }
   ],
   "source": [
    "%%px \n",
    "print(o.get_is_transposed())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] \n",
      "k0: [0 1 2 3]\n",
      "k1: [  0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15  16  17\n",
      "  18  19  20  21  22  23  24 -23 -22 -21 -20 -19 -18 -17 -16 -15 -14 -13\n",
      " -12 -11 -10  -9  -8  -7  -6  -5  -4  -3  -2  -1]\n",
      "[stdout:1] \n",
      "k0: [4 5 6 7]\n",
      "k1: [  0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15  16  17\n",
      "  18  19  20  21  22  23  24 -23 -22 -21 -20 -19 -18 -17 -16 -15 -14 -13\n",
      " -12 -11 -10  -9  -8  -7  -6  -5  -4  -3  -2  -1]\n",
      "[stdout:2] \n",
      "k0: [ 8  9 10 11]\n",
      "k1: [  0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15  16  17\n",
      "  18  19  20  21  22  23  24 -23 -22 -21 -20 -19 -18 -17 -16 -15 -14 -13\n",
      " -12 -11 -10  -9  -8  -7  -6  -5  -4  -3  -2  -1]\n",
      "[stdout:3] \n",
      "k0: [12 13 14 15]\n",
      "k1: [  0   1   2   3   4   5   6   7   8   9  10  11  12  13  14  15  16  17\n",
      "  18  19  20  21  22  23  24 -23 -22 -21 -20 -19 -18 -17 -16 -15 -14 -13\n",
      " -12 -11 -10  -9  -8  -7  -6  -5  -4  -3  -2  -1]\n"
     ]
    }
   ],
   "source": [
    "%%px \n",
    "k0, k1 = o.get_k_adim_loc()\n",
    "print('k0:', k0)\n",
    "print('k1:', k1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] \n",
      "(12, 32)\n",
      "(4, 48)\n",
      "[stdout:1] \n",
      "(12, 32)\n",
      "(4, 48)\n",
      "[stdout:2] \n",
      "(12, 32)\n",
      "(4, 48)\n",
      "[stdout:3] \n",
      "(12, 32)\n",
      "(4, 48)\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "print(o.get_shapeX_loc())\n",
    "print(o.get_shapeK_loc())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] (0, 0)\n",
      "[stdout:1] (12, 0)\n",
      "[stdout:2] (24, 0)\n",
      "[stdout:3] (36, 0)\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "print(o.get_seq_indices_first_X())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] (0, 0)\n",
      "[stdout:1] (4, 0)\n",
      "[stdout:2] (8, 0)\n",
      "[stdout:3] (12, 0)\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "print(o.get_seq_indices_first_K())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "a = np.ones(o.get_shapeX_loc())\n",
    "a_fft = o.fft(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "a_fft = np.empty(o.get_shapeK_loc(), dtype=np.complex128)\n",
    "o.fft_as_arg(a, a_fft)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "o.ifft_as_arg(a_fft, a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "a = o.ifft(a_fft)"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}