{
 "metadata": {
  "name": "pytrilinos"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from time import sleep\n",
      "import numpy as np\n",
      "from PyTrilinos import Epetra"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# initialize MPI\n",
      "Comm  = Epetra.PyComm()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "NumElements = 10\n",
      "MapUnique = Epetra.Map(NumElements, 0, Comm)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print \"Unique map\"\n",
      "print MapUnique"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Partially overlapping distribution\n",
      "LocalElements = { 0: [0, 1, 2, 3, 4, 8],\n",
      "                  1: [3, 4, 5, 6],\n",
      "                  2: [6, 7, 8, 9] \n",
      "                }\n",
      "MapOverl =  Epetra.Map(-1, LocalElements[Comm.MyPID()], 0, Comm)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print \"Partially overlapping map\"\n",
      "print MapOverl"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Importer from overlapping to unique\n",
      "Exporter = Epetra.Export(MapOverl, MapUnique)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "# Create vectors\n",
      "XOverl = Epetra.Vector(MapOverl)\n",
      "XUnique = Epetra.Vector(MapUnique) #initialized to 0"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "XOverl[:] = Comm.MyPID() + 1\n",
      "print \"Max Value\"\n",
      "print Comm.MyPID(), XOverl.MaxValue()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "Comm.Barrier()\n",
      "def printVector(label, vector):\n",
      "    Comm.Barrier()\n",
      "    if Comm.MyPID() == 0: print label\n",
      "    sleep(Comm.MyPID())\n",
      "    print (Comm.MyPID(), vector)\n",
      "printVector(\"Overlapping Vector\", XOverl)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "XUnique.Export(XOverl, Exporter, Epetra.Average)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "printVector(\"Unique Vector\", XUnique)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "M=np.array([[ 2., -1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],\n",
      "           [-1.,  2., -1.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],\n",
      "           [ 0., -1.,  2., -1.,  0.,  0.,  0.,  0.,  0.,  0.],\n",
      "           [ 0.,  0., -1.,  2., -1.,  0.,  0.,  0.,  0.,  0.],\n",
      "           [ 0.,  0.,  0., -1.,  2., -1.,  0.,  0.,  0.,  0.],\n",
      "           [ 0.,  0.,  0.,  0., -1.,  2., -1.,  0.,  0.,  0.],\n",
      "           [ 0.,  0.,  0.,  0.,  0., -1.,  2., -1.,  0.,  0.],\n",
      "           [ 0.,  0.,  0.,  0.,  0.,  0., -1.,  2., -1.,  0.],\n",
      "           [ 0.,  0.,  0.,  0.,  0.,  0.,  0., -1.,  2., -1.],\n",
      "           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0., -1.,  2.]])\n"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "MyGlobalElements  = MapUnique.MyGlobalElements()\n",
      "Matrix            = Epetra.CrsMatrix(Epetra.Copy, MapUnique, 0)\n",
      "for i in MyGlobalElements:\n",
      "    if i > 0:\n",
      "        Matrix[i, i - 1] = -1\n",
      "    if i < NumElements - 1:\n",
      "        Matrix[i, i + 1] = -1\n",
      "    Matrix[i, i] = 2.\n",
      "assert Matrix.FillComplete() == 0"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "for i in MyGlobalElements:\n",
      "    print \"%d: Matrix(%d, %d) = %e\" %(Comm.MyPID(), i, i, Matrix[i, i])"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "XUnique2 = Epetra.Vector(MapUnique)\n",
      "Matrix.Multiply(False, XUnique, XUnique2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "printVector(\"Unique Vector\", XUnique2)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from PyTrilinos import AztecOO\n",
      "Solution = Epetra.Vector(MapUnique)\n",
      "LinProb = Epetra.LinearProblem(Matrix, Solution, XUnique2)\n",
      "IterSolver = AztecOO.AztecOO(LinProb)\n",
      "IterSolver.Iterate(10, 1e-9)\n",
      "printVector(\"Correct Solution\", XUnique)\n",
      "printVector(\"Iterative Solution\", Solution)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": []
    }
   ],
   "metadata": {}
  }
 ]
}