{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5.1 Poisson Equation in Parallel\n",
    "\n",
    "NGSolve can be executed on a cluster using the MPI message passing interface.\n",
    "You can download [poisson_mpi.py](poisson_mpi.py) and run it as\n",
    "\n",
    "> mpirun -np 4 python3 poisson_mpi.py \n",
    "\n",
    "This computes and saves a solution which \n",
    "we can visualize it with [drawsolution.py](drawsolution.py)\n",
    "\n",
    "> netgen drawsolution.py\n",
    "\n",
    "For proper parallel execution, Netgen/NGSolve must be configured with\n",
    "`-DUSE_MPI=ON`. Recent binaries for Linux and Mac are built with MPI support (?). If you are unsure, when  your Netgen/NGSolve supports MPI, look for outputs like \"Including MPI version 3.1\" during Netgen startup."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## MPI-parallel execution using ipyparallel\n",
    "\n",
    "For the MPI jupyter-tutorials we use `ipyparallel` module. Please consult https://ipyparallel.readthedocs.io for installation.\n",
    "\n",
    "We start an MPI server:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ipyparallel import Cluster\n",
    "c = await Cluster(engines=\"mpi\").start_and_connect(n=4, activate=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use mpi4py https://mpi4py.readthedocs.io/ for issuing MPI calls from Python. The %%px syntax magic causes parallel execution of that cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px \n",
    "from mpi4py.MPI import COMM_WORLD as comm\n",
    "print (comm.rank, comm.size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The *master* process (rank==0) generates the mesh, and distributes it within the group of processes defined by the *communicator*. All other ranks receive a part of the mesh. The function mesh.GetNE(VOL) returns the local number of elements:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "from ngsolve import *\n",
    "from netgen.occ import unit_square\n",
    "\n",
    "if comm.rank == 0:\n",
    "    mesh = Mesh(unit_square.GenerateMesh(maxh=0.1).Distribute(comm))\n",
    "else:\n",
    "    mesh = Mesh(netgen.meshing.Mesh.Receive(comm))\n",
    "print (mesh.GetNE(VOL))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can define spaces, bilinear / linear forms, and gridfunctions in the same way as in sequential mode. But now, the degrees of freedom are distributed on the cluster following the distribution of the mesh. The finite element spaces define how the dofs match together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "fes = H1(mesh, order=3, dirichlet=\".*\")\n",
    "u,v = fes.TnT()\n",
    "\n",
    "a = BilinearForm(grad(u)*grad(v)*dx)\n",
    "pre = Preconditioner(a, \"local\")\n",
    "a.Assemble()\n",
    "\n",
    "f = LinearForm(1*v*dx).Assemble()\n",
    "gfu = GridFunction(fes)\n",
    "\n",
    "from ngsolve.krylovspace import CGSolver\n",
    "inv = CGSolver(a.mat, pre.mat, printrates=comm.rank==0, maxiter=200, tol=1e-8)\n",
    "gfu.vec.data = inv*f.vec"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Parallel pickling allows to serialize the distributed solution and transfer it to the client. The process with rank=0 gets the whole mesh and computed solution, all other processes get the local parts of the mesh and solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfu = c[:][\"gfu\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now draw the whole solution using the the master process's `gfu[0]`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve.webgui import Draw\n",
    "Draw (gfu[0]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Drawing `gfu[n]` will draw only part of the solution that the process with rank=`n` possesses. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw (gfu[3]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also visualize the sub-domains obtained by the automatic partitioning, without using any computed solution, as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "fesL2 = L2(mesh, order=0)\n",
    "gfL2 = GridFunction(fesL2)\n",
    "gfL2.vec.local_vec[:] = comm.rank"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gfL2 = c[:][\"gfL2\"]\n",
    "Draw (gfL2[0]);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}