{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting demo_2d.py\n"
     ]
    }
   ],
   "source": [
    "%%file demo_2d.py\n",
    "\n",
    "from __future__ import print_function, division\n",
    "import os\n",
    "import sys, petsc4py\n",
    "petsc4py.init(sys.argv)\n",
    "import mpi4py.MPI as mpi\n",
    "from petsc4py import PETSc\n",
    "import numpy as np\n",
    "from GenEO import *\n",
    "\n",
    "def rhs(coords, rhs):\n",
    "    n = rhs.shape\n",
    "    rhs[..., 1] = -9.81\n",
    "\n",
    "OptDB = PETSc.Options()\n",
    "Lx = OptDB.getInt('Lx', 10)\n",
    "Ly = OptDB.getInt('Ly', 1)\n",
    "n  = OptDB.getInt('n', 16)\n",
    "nx = OptDB.getInt('nx', Lx*n)\n",
    "ny = OptDB.getInt('ny', Ly*n)\n",
    "E1 = OptDB.getReal('E1', 10**6)\n",
    "E2 = OptDB.getReal('E2', 1)\n",
    "nu1 = OptDB.getReal('nu1', 0.4)\n",
    "nu2 = OptDB.getReal('nu2', 0.4)\n",
    "\n",
    "hx = Lx/(nx - 1)\n",
    "hy = Ly/(ny - 1)\n",
    "\n",
    "da = PETSc.DMDA().create([nx, ny], dof=2, stencil_width=1)\n",
    "da.setUniformCoordinates(xmax=Lx, ymax=Ly)\n",
    "da.setMatType(PETSc.Mat.Type.IS)\n",
    "da.setFieldName(0, 'u')\n",
    "da.setFieldName(1, 'v')\n",
    "\n",
    "path = './output_2d/'\n",
    "if mpi.COMM_WORLD.rank == 0:\n",
    "    if not os.path.exists(path):\n",
    "        os.mkdir(path)\n",
    "    else:\n",
    "        os.system('rm {}/*.vts'.format(path))\n",
    "\n",
    "def lame_coeff(x, y, v1, v2):\n",
    "    output = np.empty(x.shape)\n",
    "    mask = np.logical_or(np.logical_and(.2<=y, y<=.4),np.logical_and(.6<=y, y<=.8))\n",
    "    output[mask] = v1\n",
    "    output[np.logical_not(mask)] = v2\n",
    "    return output\n",
    "\n",
    "# non constant Young's modulus and Poisson's ratio \n",
    "E = buildCellArrayWithFunction(da, lame_coeff, (E1,E2))\n",
    "nu = buildCellArrayWithFunction(da, lame_coeff, (nu1,nu2))\n",
    "\n",
    "lamb = (nu*E)/((1+nu)*(1-2*nu)) \n",
    "mu = .5*E/(1+nu)\n",
    "\n",
    "class callback:\n",
    "    def __init__(self, da):\n",
    "        self.da = da\n",
    "        ranges = da.getRanges()\n",
    "        ghost_ranges = da.getGhostRanges()\n",
    "        \n",
    "        self.slices = []\n",
    "        for r, gr in zip(ranges, ghost_ranges):\n",
    "            self.slices.append(slice(gr[0], r[1]))\n",
    "        self.slices = tuple(self.slices)\n",
    "\n",
    "        self.it = 0\n",
    "\n",
    "    def __call__(self, locals):\n",
    "        pyKSP = locals['self']\n",
    "        proj = pyKSP.mpc.proj\n",
    "\n",
    "        if self.it == 0:\n",
    "            work, _ = proj.A.getVecs()\n",
    "            for i, vec in enumerate(proj.coarse_vecs):\n",
    "                if vec:\n",
    "                    proj.workl = vec.copy()\n",
    "                else:\n",
    "                    proj.workl.set(0.)\n",
    "                work.set(0)\n",
    "                proj.scatter_l2g(proj.workl, work, PETSc.InsertMode.ADD_VALUES)\n",
    "\n",
    "                viewer = PETSc.Viewer().createVTK(path + 'coarse_vec_{}.vts'.format(i), 'w', comm = PETSc.COMM_WORLD)\n",
    "                tmp = self.da.createGlobalVec()\n",
    "                tmpl_a = self.da.getVecArray(tmp)\n",
    "                work_a = self.da.getVecArray(work)\n",
    "                tmpl_a[:] = work_a[:]\n",
    "                tmp.view(viewer)\n",
    "                viewer.destroy()\n",
    "            self.it += 1\n",
    "\n",
    "\n",
    "x = da.createGlobalVec()\n",
    "b = buildRHS(da, [hx, hy], rhs)\n",
    "A = buildElasticityMatrix(da, [hx, hy], lamb, mu)\n",
    "A.assemble()\n",
    "bcApplyWest(da, A, b)\n",
    "\n",
    "#Setup the preconditioner (or multipreconditioner) and the coarse space\n",
    "pcbnn = PCBNN(A)\n",
    "\n",
    "# Set initial guess\n",
    "x.setRandom()\n",
    "xnorm = b.dot(x)/x.dot(A*x)\n",
    "x *= xnorm\n",
    "\n",
    "ksp = PETSc.KSP().create()\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(ksp.Type.PYTHON)\n",
    "pyKSP = KSP_AMPCG(pcbnn)\n",
    "pyKSP.callback = callback(da)\n",
    "ksp.setPythonContext(pyKSP)\n",
    "ksp.setFromOptions()\n",
    "ksp.setInitialGuessNonzero(True)\n",
    "\n",
    "ksp.solve(b, x)\n",
    "\n",
    "viewer = PETSc.Viewer().createVTK(path + 'solution_2d.vts', 'w', comm = PETSc.COMM_WORLD)\n",
    "x.view(viewer)\n",
    "\n",
    "\n",
    "lamb_petsc = da.createGlobalVec()\n",
    "lamb_a = da.getVecArray(lamb_petsc)\n",
    "coords = da.getCoordinates()\n",
    "coords_a = da.getVecArray(coords)\n",
    "E = lame_coeff(coords_a[..., 0], coords_a[..., 1], E1, E2)\n",
    "nu = lame_coeff(coords_a[..., 0], coords_a[..., 1], nu1, nu2)\n",
    "\n",
    "lamb_a[..., 0] = (nu*E)/((1+nu)*(1-2*nu)) \n",
    "lamb_a[..., 1] = mpi.COMM_WORLD.rank\n",
    "lamb_petsc.view(viewer)\n",
    "\n",
    "viewer.destroy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Subdomain number 0 contributes 0 coarse vectors as zero energy modes of local solver\n",
      "Subdomain number 2 contributes 3 coarse vectors as zero energy modes of local solver\n",
      "Subdomain number 1 contributes 3 coarse vectors as zero energy modes of local solver\n",
      "Subdomain number 3 contributes 3 coarse vectors as zero energy modes of local solver\n",
      "GenEO eigenvalue number 0 for lambdamax in subdomain 2: (1.4945905664998919e-06+0j)\n",
      "GenEO eigenvalue number 1 for lambdamax in subdomain 2: (1.3826304536173795e-05+0j)\n",
      "GenEO eigenvalue number 2 for lambdamax in subdomain 2: (4.5054862698830293e-05+0j)\n",
      "GenEO eigenvalue number 3 for lambdamax in subdomain 2: (0.05405868128473588+0j)\n",
      "GenEO eigenvalue number 4 for lambdamax in subdomain 2: (0.054059009027751584+0j)\n",
      "GenEO eigenvalue number 5 for lambdamax in subdomain 2: (0.05666304812554981+0j)\n",
      "GenEO eigenvalue number 6 for lambdamax in subdomain 2: (0.05707962286854576+0j)\n",
      "GenEO eigenvalue number 7 for lambdamax in subdomain 2: (0.15901212849796265+0j)\n",
      "GenEO eigenvalue number 8 for lambdamax in subdomain 2: (0.15905866059814386+0j)\n",
      "GenEO eigenvalue number 9 for lambdamax in subdomain 2: (0.22127754627469237+0j) <-- not selected (> 0.2)\n",
      "Subdomain number 2 contributes 12 coarse vectors after first GenEO\n",
      "Subdomain number 2 contributes 12 coarse vectors in total\n",
      "GenEO eigenvalue number 0 for lambdamax in subdomain 1: (1.4945905664998919e-06+0j)\n",
      "GenEO eigenvalue number 1 for lambdamax in subdomain 1: (1.3826304536173795e-05+0j)\n",
      "GenEO eigenvalue number 2 for lambdamax in subdomain 1: (4.5054862698830293e-05+0j)\n",
      "GenEO eigenvalue number 3 for lambdamax in subdomain 1: (0.05405868128473588+0j)\n",
      "GenEO eigenvalue number 4 for lambdamax in subdomain 1: (0.054059009027751584+0j)\n",
      "GenEO eigenvalue number 5 for lambdamax in subdomain 1: (0.05666304812554981+0j)\n",
      "GenEO eigenvalue number 6 for lambdamax in subdomain 1: (0.05707962286854576+0j)\n",
      "GenEO eigenvalue number 7 for lambdamax in subdomain 1: (0.15901212849796265+0j)\n",
      "GenEO eigenvalue number 8 for lambdamax in subdomain 1: (0.15905866059814386+0j)\n",
      "GenEO eigenvalue number 9 for lambdamax in subdomain 1: (0.22127754627469237+0j) <-- not selected (> 0.2)\n",
      "Subdomain number 1 contributes 12 coarse vectors after first GenEO\n",
      "Subdomain number 1 contributes 12 coarse vectors in total\n",
      "GenEO eigenvalue number 0 for lambdamax in subdomain 0: (0.0003173241987885001+0j)\n",
      "GenEO eigenvalue number 1 for lambdamax in subdomain 0: (0.0003370240261594047+0j)\n",
      "GenEO eigenvalue number 2 for lambdamax in subdomain 0: (0.02828240156782736+0j)\n",
      "GenEO eigenvalue number 3 for lambdamax in subdomain 0: (0.028283168511413163+0j)\n",
      "GenEO eigenvalue number 4 for lambdamax in subdomain 0: (0.11198298627142454+0j)\n",
      "GenEO eigenvalue number 5 for lambdamax in subdomain 0: (0.11219184197639594+0j)\n",
      "GenEO eigenvalue number 6 for lambdamax in subdomain 0: (0.2212907727033022+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 7 for lambdamax in subdomain 0: (0.22129077350512905+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 8 for lambdamax in subdomain 0: (0.5564840670455413+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 9 for lambdamax in subdomain 0: (0.5839813243185217+0j) <-- not selected (> 0.2)\n",
      "Subdomain number 0 contributes 6 coarse vectors after first GenEO\n",
      "This is BNN so eigmin = 1, no eigenvalue problem will be solved for eigmin\n",
      "Subdomain number 0 contributes 6 coarse vectors in total\n",
      "GenEO eigenvalue number 0 for lambdamax in subdomain 3: (2.120823483577397e-05+0j)\n",
      "GenEO eigenvalue number 1 for lambdamax in subdomain 3: (3.185852816809119e-05+0j)\n",
      "GenEO eigenvalue number 2 for lambdamax in subdomain 3: (0.008267951049446151+0j)\n",
      "GenEO eigenvalue number 3 for lambdamax in subdomain 3: (0.22129069514149075+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 4 for lambdamax in subdomain 3: (0.22129069789192463+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 5 for lambdamax in subdomain 3: (0.5564838682522086+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 6 for lambdamax in subdomain 3: (0.5839813228930885+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 7 for lambdamax in subdomain 3: (0.583981322910434+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 8 for lambdamax in subdomain 3: (0.6574284065235129+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 9 for lambdamax in subdomain 3: (0.6574287754614379+0j) <-- not selected (> 0.2)\n",
      "GenEO eigenvalue number 10 for lambdamax in subdomain 3: (0.8078884348443373+0j) <-- not selected (> 0.2)\n",
      "Subdomain number 3 contributes 6 coarse vectors after first GenEO\n",
      "Subdomain number 3 contributes 6 coarse vectors in total\n",
      "There are 36 vectors in the coarse space.\n",
      "multipreconditioning initial iteration\n",
      "  0 KSP Residual norm 2.064529884745e+02 \n",
      "\tnatural_norm -> 1.73556393e+00\n",
      "\tti -> 0.00000000e+00\n",
      "  1 KSP Residual norm 3.932753275825e+01 \n",
      "\tnatural_norm -> 1.63391933e+00\n",
      "\tti -> 9.08271064e-01\n",
      "  2 KSP Residual norm 1.580952241755e+01 \n",
      "\tnatural_norm -> 2.18092684e-01\n",
      "\tti -> 9.35102630e+00\n",
      "  3 KSP Residual norm 9.120636856569e+00 \n",
      "\tnatural_norm -> 1.31680828e-01\n",
      "\tti -> 9.19892997e-01\n",
      "  4 KSP Residual norm 5.183587708442e+00 \n",
      "\tnatural_norm -> 1.39142872e-02\n",
      "\tti -> 1.66585316e+01\n",
      "  5 KSP Residual norm 4.620210219752e-01 \n",
      "\tnatural_norm -> 3.69402830e-03\n",
      "\tti -> 9.83343825e+00\n",
      "  6 KSP Residual norm 3.409145943303e-01 \n",
      "\tnatural_norm -> 9.44409737e-04\n",
      "\tti -> 2.51502577e+00\n",
      "  7 KSP Residual norm 6.176398129082e-02 \n",
      "\tnatural_norm -> 1.94856155e-04\n",
      "\tti -> 1.41290595e+01\n",
      "  8 KSP Residual norm 5.612980151526e-03 \n",
      "\tnatural_norm -> 2.96173613e-05\n",
      "\tti -> 2.49126244e+01\n",
      "  9 KSP Residual norm 6.261411769011e-04 \n",
      "\tnatural_norm -> 8.52664601e-06\n",
      "\tti -> 7.32377097e+00\n",
      " 10 KSP Residual norm 4.518137976821e-04 \n",
      "\tnatural_norm -> 4.25968757e-07\n",
      "\tti -> 2.62089341e+02\n",
      " 11 KSP Residual norm 1.318225537479e-04 \n",
      "\tnatural_norm -> 4.12227280e-07\n",
      "\tti -> 6.07461529e-01\n",
      " 12 KSP Residual norm 2.998427485754e-05 \n",
      "\tnatural_norm -> 4.89607585e-08\n",
      "\tti -> 1.29464708e+01\n",
      " 13 KSP Residual norm 8.616862952616e-06 \n",
      "\tnatural_norm -> 6.81579981e-09\n",
      "\tti -> 2.65471292e+01\n"
     ]
    }
   ],
   "source": [
    "!mpiexec -np 4 python demo_2d.py -AMPCG_verbose -ksp_monitor -PCBNN_verbose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from plot import plot_solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aff8a370337240e29752c938ad3505f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/html": [
       "<p>Failed to display Jupyter Widget of type <code>HBox</code>.</p>\n",
       "<p>\n",
       "  If you're reading this message in the Jupyter Notebook or JupyterLab Notebook, it may mean\n",
       "  that the widgets JavaScript is still loading. If this message persists, it\n",
       "  likely means that the widgets JavaScript library is either not installed or\n",
       "  not enabled. See the <a href=\"https://ipywidgets.readthedocs.io/en/stable/user_install.html\">Jupyter\n",
       "  Widgets Documentation</a> for setup instructions.\n",
       "</p>\n",
       "<p>\n",
       "  If you're reading this message in another frontend (for example, a static\n",
       "  rendering on GitHub or <a href=\"https://nbviewer.jupyter.org/\">NBViewer</a>),\n",
       "  it may mean that your frontend doesn't currently support widgets.\n",
       "</p>\n"
      ],
      "text/plain": [
       "HBox(children=(Renderer(camera=PerspectiveCamera(position=(0.0, 0.0, 13.0), quaternion=(0.0, 0.0, 0.0, 1.0), scale=(1.0, 1.0, 1.0), up=(0.0, 1.0, 0.0)), controls=[OrbitControls(controlling=PerspectiveCamera(position=(0.0, 0.0, 13.0), quaternion=(0.0, 0.0, 0.0, 1.0), scale=(1.0, 1.0, 1.0), up=(0.0, 1.0, 0.0)), enableRotate=False)], scene=Scene(background='black', children=(Mesh(geometry=BufferGeometry(attributes={'position': <BufferAttribute shape=(14310, 3), dtype=float32>, 'color': <BufferAttribute shape=(4770, 3, 3), dtype=float32>}), material=MeshLambertMaterial(alphaMap=None, aoMap=None, color='0xF5F5F5', emissiveMap=None, envMap=None, lightMap=None, map=None, opacity=0.5, side='DoubleSide', specularMap=None, transparent=True, vertexColors='VertexColors'), quaternion=(0.0, 0.0, 0.0, 1.0), scale=(1.0, 1.0, 1.0), up=(0.0, 1.0, 0.0)), LineSegments(geometry=BufferGeometry(attributes={'position': <BufferAttribute shape=(19080, 3), dtype=float32>}), material=LineBasicMaterial(color='0xffffff'), quaternion=(0.0, 0.0, 0.0, 1.0), scale=(1.0, 1.0, 1.0), up=(0.0, 1.0, 0.0)), AmbientLight(quaternion=(0.0, 0.0, 0.0, 1.0), scale=(1.0, 1.0, 1.0), up=(0.0, 1.0, 0.0))), fog=None, overrideMaterial=None, position=(-5.192105760650471, 2.9941447310517106, -0.0), quaternion=(0.0, 0.0, 0.0, 1.0), scale=(1.0, 1.0, 1.0), up=(0.0, 1.0, 0.0)), shadowMap=WebGLShadowMap()), VBox(children=(Dropdown(description='Fields', options=OrderedDict([('u', 0), ('v', 1), ('lambda', 2), ('rank', 3)]), value=0), Dropdown(description='domains', options=OrderedDict([('all', -1), ('domain 0', 0), ('domain 1', 1), ('domain 2', 2), ('domain 3', 3)]), value=-1), Dropdown(description='coarse vecs', options=OrderedDict([('none', -1), ('coarse vec 0', 0), ('coarse vec 1', 1), ('coarse vec 2', 2), ('coarse vec 3', 3), ('coarse vec 4', 4), ('coarse vec 5', 5), ('coarse vec 6', 6), ('coarse vec 7', 7), ('coarse vec 8', 8), ('coarse vec 9', 9), ('coarse vec 10', 10), ('coarse vec 11', 11), ('coarse vec 12', 12), ('coarse vec 13', 13), ('coarse vec 14', 14), ('coarse vec 15', 15), ('coarse vec 16', 16), ('coarse vec 17', 17), ('coarse vec 18', 18), ('coarse vec 19', 19), ('coarse vec 20', 20), ('coarse vec 21', 21), ('coarse vec 22', 22), ('coarse vec 23', 23), ('coarse vec 24', 24), ('coarse vec 25', 25), ('coarse vec 26', 26), ('coarse vec 27', 27), ('coarse vec 28', 28), ('coarse vec 29', 29), ('coarse vec 30', 30), ('coarse vec 31', 31), ('coarse vec 32', 32), ('coarse vec 33', 33), ('coarse vec 34', 34), ('coarse vec 35', 35)]), value=-1), Checkbox(value=True, description='Show displacement'), Checkbox(value=True, description='Show mesh'), IntSlider(value=100, continuous_update=False, description='Scale factor for coarse vec displacement', max=1000, min=1)))))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_solution('output_2d', 'solution_2d.vts')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}