{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5.2 Parallel dofs and Vector-types\n",
    "In this tutorial we learn how NGSolve represents distributed finite element spaces."
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "from ngsolve import *\n",
    "from netgen.geom2d import unit_square\n",
    "comm = MPI.COMM_WORLD\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))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create the space on the distributed mesh. All processes agree on the global number of dofs. Locally, each rank has access only to the subset of dofs associated with its elements. Some dofs are shared by several ranks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "fes = H1(mesh, order=1)\n",
    "# fes = L2(mesh, order=0)\n",
    "print (\"global dofs =\", fes.ndofglobal, \", local dofs =\", fes.ndof, \\\n",
    "       \", sum of local dofs =\", comm.allreduce(fes.ndof))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parallel Dofs\n",
    "A ParallelDofs object maintains information how dofs are connected across the cluster. The ParallelDofs object is generated by the FESpace, which has access to the connectivity of the mesh."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "pardofs = fes.ParallelDofs()\n",
    "for k in range(pardofs.ndoflocal):\n",
    "    print (\"dof\", k, \"is shard with ranks\", list(pardofs.Dof2Proc(k)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "print (\"I share dofs with ranks:\", list(pardofs.ExchangeProcs()))\n",
    "for k in range(MPI.COMM_WORLD.size):\n",
    "    print (\"with rank\", k, \"I share dofs\", list(pardofs.Proc2Dof(k)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "u,v = fes.TnT()\n",
    "M = BilinearForm(u*v*dx).Assemble().mat\n",
    "gfu = GridFunction(fes)\n",
    "gfu.Set (1)\n",
    "print (gfu.vec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that all values are set to 1, i.e. joint dofs have the same value. We call such a vector 'cumulated'. The matrix M is stored locally assembled, i.e. every rank has the contributions from its elements. When we multiply this matrix with a cumulated vector, every rank performs a local matrix vector product. The resulting vector is stored 'distributed', i.e. the true values are obtained by adding up rank-local contributions for joint dofs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "r = M.CreateColVector()\n",
    "r.data = M*gfu.vec\n",
    "print (r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This cumulated/distributed pair of vectors is prefect for computing inner products. We can compute inner products of local vectors, and sum up (i.e. reduce in MPI terms) across all ranks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "print (\"global ip =\", InnerProduct(gfu.vec, r))\n",
    "localip = InnerProduct(r.local_vec, gfu.vec.local_vec)\n",
    "print (\"local contribution:\", localip)\n",
    "print (\"cumulated:\", comm.allreduce(localip, MPI.SUM))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}