{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial FFT 3D parallel (MPI)\n",
    "\n",
    "In this tutorial, we present how to use fluidfft to perform 3D fft in sequential.\n",
    "\n",
    "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": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "import numpy as np\n",
    "from fluidfft.fft3d import methods_mpi\n",
    "from fluidfft import import_fft_class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['fft3d.mpi_with_fftw1d', 'fft3d.mpi_with_fftwmpi3d', 'fft3d.mpi_with_p3dfft', 'fft3d.mpi_with_pfft']\n"
     ]
    }
   ],
   "source": [
    "%%px --targets 1\n",
    "print(methods_mpi)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We import a class and instantiate it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "cls = import_fft_class('fft3d.mpi_with_fftw1d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "o = cls(4, 8, 16)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's have a look at the attribute of this objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "build_invariant_arrayK_from_2d_indices12X\n",
      "build_invariant_arrayX_from_2d_indices12X\n",
      "comm\n",
      "compute_energy_from_K\n",
      "compute_energy_from_X\n",
      "create_arrayK\n",
      "create_arrayX\n",
      "fft\n",
      "fft_as_arg\n",
      "gather_Xspace\n",
      "get_dimX_K\n",
      "get_dim_first_fft\n",
      "get_k_adim_loc\n",
      "get_local_size_X\n",
      "get_seq_indices_first_K\n",
      "get_seq_indices_first_X\n",
      "get_shapeK_loc\n",
      "get_shapeK_seq\n",
      "get_shapeX_loc\n",
      "get_shapeX_seq\n",
      "get_short_name\n",
      "ifft\n",
      "ifft_as_arg\n",
      "ifft_as_arg_destroy\n",
      "nb_proc\n",
      "rank\n",
      "run_benchs\n",
      "run_tests\n",
      "scatter_Xspace\n",
      "sum_wavenumbers\n"
     ]
    }
   ],
   "source": [
    "%%px --targets 1\n",
    "print('\\n'.join([name for name in dir(o) if not name.startswith('__')]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's run a test and benchmark the fft and ifft functions directly from C++."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "_ = o.run_tests()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] t_fft = 0.0029698 s; t_ifft = 4.11e-05 s\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "results = o.run_benchs()\n",
    "if rank == 0:\n",
    "    print('t_fft = {} s; t_ifft = {} s'.format(*results))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's understand how the data is stored:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] (2, 1, 0)\n",
      "[stdout:1] (2, 1, 0)\n",
      "[stdout:2] (2, 1, 0)\n",
      "[stdout:3] (2, 1, 0)\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "print(o.get_dimX_K())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which means that for this class, in Fourier space, the data is transposed...\n",
    "\n",
    "Now we can get the non dimensional wavenumber in the first and second dimensions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] \n",
      "k0: [0 1]\n",
      "k1: [ 0  1  2  3  4 -3 -2 -1]\n",
      "k2: [ 0  1  2 -1]\n",
      "[stdout:1] \n",
      "k0: [2 3]\n",
      "k1: [ 0  1  2  3  4 -3 -2 -1]\n",
      "k2: [ 0  1  2 -1]\n",
      "[stdout:2] \n",
      "k0: [4 5]\n",
      "k1: [ 0  1  2  3  4 -3 -2 -1]\n",
      "k2: [ 0  1  2 -1]\n",
      "[stdout:3] \n",
      "k0: [6 7]\n",
      "k1: [ 0  1  2  3  4 -3 -2 -1]\n",
      "k2: [ 0  1  2 -1]\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "k0, k1, k2 = o.get_k_adim_loc()\n",
    "print('k0:', k0)\n",
    "print('k1:', k1)\n",
    "print('k2:', k2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] (0, 0, 0)\n",
      "[stdout:1] (2, 0, 0)\n",
      "[stdout:2] (4, 0, 0)\n",
      "[stdout:3] (6, 0, 0)\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "print(o.get_seq_indices_first_K())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and get the shape of the arrays in real and Fourier space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] (4, 8, 16) (1, 8, 16)\n",
      "[stdout:1] (4, 8, 16) (1, 8, 16)\n",
      "[stdout:2] (4, 8, 16) (1, 8, 16)\n",
      "[stdout:3] (4, 8, 16) (1, 8, 16)\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "print(o.get_shapeX_seq(), o.get_shapeX_loc())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[stdout:0] (8, 8, 4) (2, 8, 4)\n",
      "[stdout:1] (8, 8, 4) (2, 8, 4)\n",
      "[stdout:2] (8, 8, 4) (2, 8, 4)\n",
      "[stdout:3] (8, 8, 4) (2, 8, 4)\n"
     ]
    }
   ],
   "source": [
    "%%px\n",
    "print(o.get_shapeK_seq(), o.get_shapeK_loc())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's compute fast Fourier transforms. We first initialize arrays:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "a = np.ones(o.get_shapeX_loc())\n",
    "a_fft = np.empty(o.get_shapeK_loc(), dtype=np.complex128)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we do not have the array where to put the result we can do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "a_fft = o.fft(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we already have the array where to put the result we can do:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "o.fft_as_arg(a, a_fft)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally for the inverse Fourier transform:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "a = o.ifft(a_fft)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "o.ifft_as_arg(a_fft, a)"
   ]
  }
 ],
 "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
}